From 34202f3f50297e9931de90459ac178d5fffcfc6d Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Fri, 14 Aug 2026 13:15:05 +0800 Subject: [PATCH 01/21] feat(cpu): add LFM2.5-2.6B support --- CMakeLists.txt | 8 + README.md | 1 + examples/CMakeLists.txt | 1 + examples/lfm2/CMakeLists.txt | 3 + examples/lfm2/README.md | 119 ++++++ examples/lfm2/benchmark_harness.hpp | 118 ++++++ examples/lfm2/config_2.6B_w4a32_kai.json | 30 ++ examples/lfm2/demo_prompt.txt | 1 + examples/lfm2/main.cpp | 274 +++++++++++++ examples/lfm2/quant_cfg_2.6B_w4a32_kai.json | 26 ++ examples/lfm2/test_validators.py | 42 ++ examples/lfm2/validate_checkpoint.py | 199 +++++++++ examples/lfm2/validate_converted_model.py | 125 ++++++ mllm/backends/cpu/CMakeLists.txt | 11 +- mllm/backends/cpu/kernels/arm/linear/kai.cpp | 52 +++ mllm/backends/cpu/kernels/arm/linear/kai.hpp | 11 + .../kernels/common/gdn/gated_delta_net.cpp | 56 +++ .../kernels/common/gdn/gated_delta_net.hpp | 6 + mllm/backends/cpu/ops/LinearOp.cpp | 109 ++++- mllm/backends/cpu/ops/LinearOp.hpp | 14 + mllm/models/lfm2/configuration_lfm2.hpp | 168 ++++++++ mllm/models/lfm2/modeling_lfm2.hpp | 387 ++++++++++++++++++ mllm/models/lfm2/tokenization_lfm2.hpp | 298 ++++++++++++++ .../llm_components/GroupedQueryAttention.hpp | 136 ++++-- tests/cpu/CMakeLists.txt | 14 + tests/cpu/Lfm2ConfigTest.cpp | 55 +++ tests/cpu/Lfm2ShortConvTest.cpp | 73 ++++ tests/cpu/Lfm2TokenizerTest.cpp | 62 +++ tests/cpu/Qwen35GDNConvTest.cpp | 54 +++ tests/nn/GroupedQueryAttentionTest.cpp | 37 ++ 30 files changed, 2464 insertions(+), 26 deletions(-) create mode 100644 examples/lfm2/CMakeLists.txt create mode 100644 examples/lfm2/README.md create mode 100644 examples/lfm2/benchmark_harness.hpp create mode 100644 examples/lfm2/config_2.6B_w4a32_kai.json create mode 100644 examples/lfm2/demo_prompt.txt create mode 100644 examples/lfm2/main.cpp create mode 100644 examples/lfm2/quant_cfg_2.6B_w4a32_kai.json create mode 100644 examples/lfm2/test_validators.py create mode 100644 examples/lfm2/validate_checkpoint.py create mode 100644 examples/lfm2/validate_converted_model.py create mode 100644 mllm/models/lfm2/configuration_lfm2.hpp create mode 100644 mllm/models/lfm2/modeling_lfm2.hpp create mode 100644 mllm/models/lfm2/tokenization_lfm2.hpp create mode 100644 tests/cpu/Lfm2ConfigTest.cpp create mode 100644 tests/cpu/Lfm2ShortConvTest.cpp create mode 100644 tests/cpu/Lfm2TokenizerTest.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index b16ae832a..1d7d2f1f2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -41,6 +41,14 @@ option(MLLM_BLAS_VENDOR_BLIS "Enable BLIS BLAS for multi-platform" OFF) # CPU Backend: SME2 and SVE2 option(MLLM_CPU_BACKEND_USE_SME2 "Enable SME2" OFF) +option( + MLLM_ARM_CPU_BACKEND_USE_OPENMP + "Compile ARM CPU backend operators and kernels with OpenMP" + ON) +option( + MLLM_ARM_KAI_USE_OPENMP + "Compile only the ARM KleidiAI linear-kernel translation unit with OpenMP" + ON) # Ascend Backend: Options option(MLLM_ASCEND_CPU_DEBUG_MODE "Enable CPU Debug mode in ascend" OFF) 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/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..619ef219e --- /dev/null +++ b/examples/lfm2/README.md @@ -0,0 +1,119 @@ +# 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, keep runtime OpenMP enabled but configure the CPU +backend without backend-wide OpenMP. LFM2.5 still selects the shared W4A32 +I8MM prefill path and retained KAI decode workspace; avoiding OpenMP regions in +every backend operator preserves decode latency for its 8-attention / 22-conv +hybrid schedule. + +```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_ARM_CPU_BACKEND_USE_OPENMP=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. + +## 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..e06af55e3 --- /dev/null +++ b/examples/lfm2/test_validators.py @@ -0,0 +1,42 @@ +# 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_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..69e2f3e7d --- /dev/null +++ b/examples/lfm2/validate_checkpoint.py @@ -0,0 +1,199 @@ +# 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"}, +} +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_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()) + 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_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..24c0a44c8 100644 --- a/mllm/backends/cpu/CMakeLists.txt +++ b/mllm/backends/cpu/CMakeLists.txt @@ -150,10 +150,19 @@ if(MLLM_KERNEL_USE_THREADS AND MLLM_KERNEL_THREADS_VENDOR_OPENMP) # 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) + if(MLLM_BUILD_ARM_BACKEND AND MLLM_ARM_CPU_BACKEND_USE_OPENMP) 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}) + elseif(MLLM_BUILD_ARM_BACKEND AND MLLM_ARM_KAI_USE_OPENMP) + # The KAI linear helper owns its tile-parallel loops. Compiling only + # this translation unit with OpenMP lets quantized Linear reuse those + # loops without adding parallel-region overhead to every CPU operator. + set_source_files_properties( + ${CMAKE_CURRENT_LIST_DIR}/kernels/arm/linear/kai.cpp + PROPERTIES COMPILE_OPTIONS "${OpenMP_CXX_FLAGS}") + target_link_libraries(MllmCPUBackend PUBLIC ${OpenMP_CXX_FLAGS}) + target_include_directories(MllmCPUBackend PUBLIC ${OpenMP_CXX_INCLUDE_DIR}) endif() endif() endif() diff --git a/mllm/backends/cpu/kernels/arm/linear/kai.cpp b/mllm/backends/cpu/kernels/arm/linear/kai.cpp index 1a7d40baf..a5bf15689 100644 --- a/mllm/backends/cpu/kernels/arm/linear/kai.cpp +++ b/mllm/backends/cpu/kernels/arm/linear/kai.cpp @@ -445,6 +445,58 @@ void KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::matmul(float* __restrict__ dst, con } } +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) { + if (lhs_fp32 == nullptr || projections == nullptr || projection_count < 2 || workspace == nullptr || 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(1, K, ukernel.get_mr(), ukernel.get_kr(), ukernel.get_sr(), 0, lhs_fp32, + K * sizeof(float), workspace); + + const size_t n_step = static_cast(ukernel.get_n_step()); + size_t total_tiles = 0; + for (size_t projection_index = 0; projection_index < projection_count; ++projection_index) { + total_tiles += (static_cast(projections[projection_index].n) + n_step - 1) / n_step; + } + const void* lhs_ptr = + static_cast(static_cast(workspace) + ukernel.get_lhs_packed_offset(0, K)); + + 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_tiles = (static_cast(projections[projection_index].n) + n_step - 1) / n_step; + if (local_tile < projection_tiles) { break; } + local_tile -= projection_tiles; + } + + if (projection_index < projection_count) { + const auto& projection = projections[projection_index]; + const int n_index = static_cast(local_tile * n_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* 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(0, n_index, dst_stride)); + + ukernel.run_matmul(1, actual_n, K, 32, lhs_ptr, rhs_ptr, dst_ptr, dst_stride, sizeof(float), + -std::numeric_limits::max(), std::numeric_limits::max()); + } + }); + return true; +} + 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..e282bcc5c 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,10 @@ 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_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/common/gdn/gated_delta_net.cpp b/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.cpp index b20e6fc3b..db279cf26 100644 --- a/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.cpp +++ b/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.cpp @@ -243,6 +243,62 @@ void depthwiseCausalConvF32(const float* input, const float* weight, float* stat } } +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__) + // LFM2.5 uses K=3. Four adjacent channels are deinterleaved 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]; + } + } + } +} + void gatedDeltaRuleF32(const float* q, const float* k, const float* v, const float* a, const float* b, const float* a_log, const float* dt_bias, float* state, float* output, int batch_size, int sequence_length, int num_key_heads, int num_value_heads, int key_head_dim, int value_head_dim) { diff --git a/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp b/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp index 78a53ac81..e95171e42 100644 --- a/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp +++ b/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp @@ -10,6 +10,12 @@ namespace mllm::cpu::gdn { void depthwiseCausalConvF32(const float* input, const float* weight, float* state, float* output, int batch_size, int sequence_length, int channels, int kernel_size); +// Same [B, S, C] / [B, C, K - 1] state contract, with the accumulation order +// used by CPUConv1D: zero, historical taps in ascending order, then current. +// Keeping this explicit lets callers retain bitwise-sensitive model semantics. +void depthwiseCausalConvHistoryFirstF32(const float* input, const float* weight, float* state, float* output, + int batch_size, int sequence_length, int channels, int kernel_size); + // Gated-delta recurrence. Each [batch, value_head] state is independent and // may run in parallel, while tokens within one state remain strictly ordered. // diff --git a/mllm/backends/cpu/ops/LinearOp.cpp b/mllm/backends/cpu/ops/LinearOp.cpp index 7d93754a1..8309acd75 100644 --- a/mllm/backends/cpu/ops/LinearOp.cpp +++ b/mllm/backends/cpu/ops/LinearOp.cpp @@ -1,10 +1,13 @@ // Copyright (c) MLLM Team. // Licensed under the MIT License. +#include #include #include #include #include +#include +#include #if defined(__linux__) #include @@ -14,6 +17,7 @@ #include "mllm/backends/cpu/kernels/Kernels.hpp" #include "mllm/core/DataTypes.hpp" #include "mllm/core/aops/LinearOp.hpp" +#include "mllm/engine/Context.hpp" namespace mllm::cpu { @@ -72,6 +76,19 @@ void traceKaiW4A32PrefillTile(KaiW4A32Tile tile, int m, int k, int n, int thread CPULinearOp::CPULinearOp(const aops::LinearOpOptions& options) : LinearOp(options) {} +void CPULinearOp::setKaiW4A32ThreadCaps(int decode_thread_cap, int prefill_thread_cap) { + if (decode_thread_cap <= 0 || prefill_thread_cap <= 0) { + throw std::invalid_argument("KAI W4A32 thread caps must be positive"); + } + kai_w4a32_decode_thread_cap_ = decode_thread_cap; + kai_w4a32_prefill_thread_cap_ = prefill_thread_cap; +} + +int CPULinearOp::kaiW4A32ThreadCount(int m) const { + return detail::kaiW4A32ThreadCount( + m, options_.getThreads(), kai_w4a32_decode_thread_cap_, 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(); } @@ -81,6 +98,93 @@ Tensor CPULinearOp::acquireKaiWorkspace(int32_t workspace_size, int m) { return kai_decode_workspace_; } +bool CPULinearOp::tryForwardSharedInputKaiM1(const Tensor& input, const BaseOp::ptr_t* linear_ops, size_t linear_op_count, + std::vector& outputs) { +#if defined(MLLM_HOST_ARCH_ARM64) || defined(MLLM_HOST_ARCH_ARM) + constexpr size_t kMaximumSharedProjections = 3; + constexpr auto kRequiredImpl = aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32; + using KaiHelper = ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk; + constexpr auto kTile = KaiHelper::Tiles::qai8dxp1x8_qsi4c32p8x8_1x8x32; + + if (Context::instance().thisThread()->trace_mode || input.isNil() || input.device() != kCPU || input.dtype() != kFloat32 + || !input.isContiguous() || input.rank() < 2 || input.size(-2) != 1 || input.size(-1) <= 0 || !outputs.empty() + || linear_ops == nullptr || linear_op_count < 2 || linear_op_count > kMaximumSharedProjections) { + return false; + } + + const auto input_shape = input.shape(); + for (size_t index = 0; index + 2 < input_shape.size(); ++index) { + if (input_shape[index] != 1) { return false; } + } + + const int32_t K = input.size(-1); + int32_t thread_count = 0; + std::array ops{}; + for (size_t index = 0; index < linear_op_count; ++index) { + auto* op = dynamic_cast(linear_ops[index].get()); + if (op == nullptr || op->getDevice() != kCPU || op->options_.impl_type != kRequiredImpl || op->options_.bias + || op->options_.in_channels != K || op->options_.out_channels <= 0 || op->weight_.isNil() + || op->weight_.device() != kCPU || op->options_.getThreads() <= 0) { + return false; + } + if (index == 0) { + thread_count = op->kaiW4A32ThreadCount(1); + } else if (op->kaiW4A32ThreadCount(1) != thread_count) { + return false; + } + ops[index] = op; + } + + std::vector prepared_outputs; + prepared_outputs.reserve(linear_op_count); + for (size_t index = 0; index < linear_op_count; ++index) { + auto output_shape = input_shape; + output_shape.back() = ops[index]->options_.out_channels; + prepared_outputs.emplace_back(Tensor::empty(output_shape, kFloat32, kCPU).alloc()); + } + + std::array projections{}; + for (size_t index = 0; index < linear_op_count; ++index) { + projections[index] = { + .dst = prepared_outputs[index].ptr(), + .packed_weight_bias = reinterpret_cast(ops[index]->weight_.ptr()), + .n = ops[index]->options_.out_channels, + }; + } + + KaiHelper kai_helper; + const size_t workspace_size = kai_helper.workspace_size(1, K, kTile); + if (workspace_size == 0 || workspace_size > static_cast(std::numeric_limits::max())) { return false; } + auto workspace = ops[0]->acquireKaiWorkspace(static_cast(workspace_size), 1); + if (!kai_helper.matmul_shared_input_m1(input.ptr(), projections.data(), linear_op_count, workspace.ptr(), + K, kTile, thread_count)) { + return false; + } + + static const bool trace_activation = [] { + const char* value = std::getenv("MLLM_KAI_SHARED_INPUT_TRACE"); + return value != nullptr && value[0] == '1' && value[1] == '\0'; + }(); + if (trace_activation) { + const uint32_t activation_bit = 1U << static_cast(linear_op_count - 2); + 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 k=%d threads=%d\n", linear_op_count, K, thread_count); + } + } + + outputs = std::move(prepared_outputs); + return true; +#else + (void)input; + (void)linear_ops; + (void)linear_op_count; + (void)outputs; + return false; +#endif +} + void CPULinearOp::load(const ParameterFile::ptr_t& ploader) { switch (ploader->version()) { case ModelFileVersion::kV1: { @@ -262,7 +366,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 +375,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: { diff --git a/mllm/backends/cpu/ops/LinearOp.hpp b/mllm/backends/cpu/ops/LinearOp.hpp index e189530ee..077dc1e7b 100644 --- a/mllm/backends/cpu/ops/LinearOp.hpp +++ b/mllm/backends/cpu/ops/LinearOp.hpp @@ -13,6 +13,11 @@ constexpr bool shouldUseKaiW4A32I8mmPrefill(int m, bool disabled, bool cpu_suppo return m >= 4 && !disabled && cpu_supports_i8mm; } +constexpr int kaiW4A32ThreadCount(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 detail class CPULinearOp final : public aops::LinearOp { @@ -25,10 +30,19 @@ class CPULinearOp final : public aops::LinearOp { void reshape(const std::vector& inputs, std::vector& outputs) override; + static bool tryForwardSharedInputKaiM1(const Tensor& input, const BaseOp::ptr_t* linear_ops, size_t linear_op_count, + std::vector& outputs); + + void setKaiW4A32ThreadCaps(int decode_thread_cap, int prefill_thread_cap); + private: Tensor acquireKaiWorkspace(int32_t workspace_size, int m); + [[nodiscard]] int kaiW4A32ThreadCount(int m) const; + Tensor kai_decode_workspace_; + int kai_w4a32_decode_thread_cap_ = 0; + int kai_w4a32_prefill_thread_cap_ = 0; }; class CPULinearOpFactory : public TypedOpFactory { diff --git a/mllm/models/lfm2/configuration_lfm2.hpp b/mllm/models/lfm2/configuration_lfm2.hpp new file mode 100644 index 000000000..383d85b23 --- /dev/null +++ b/mllm/models/lfm2/configuration_lfm2.hpp @@ -0,0 +1,168 @@ +// 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"); + 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..b8fabb0d6 --- /dev/null +++ b/mllm/models/lfm2/modeling_lfm2.hpp @@ -0,0 +1,387 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mllm/backends/cpu/ops/LinearOp.hpp" +#include "mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp" +#include "mllm/core/Tensor.hpp" +#include "mllm/models/ARGeneration.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/llm_components/GroupedQueryAttention.hpp" +#include "mllm/nn/lmcache/KVHeadStaticCache.hpp" + +namespace mllm::models::lfm2 { + +inline void configureLfm2KaiW4A32Threads(nn::Linear& linear) { + auto op = std::dynamic_pointer_cast(linear.impl()->getInstancedOp()); + if (op != nullptr) { + // OnePlus 13T source-bound screening keeps I8MM prefill above 80 tok/s + // with six workers, while four workers avoid decode GEMV oversubscription. + op->setKaiW4A32ThreadCaps(4, 6); + } +} + +// Model-level orchestration: materialize the immutable analytical RoPE table. +// Rotation itself remains the registered nn::RoPE operation. +inline auto makeRoPEInvFreq(int32_t head_dim, float rope_theta) -> Tensor { + 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; +} + +inline auto makeRotaryPosEmbedding(const Tensor& position_ids, const Tensor& inv_freq) -> std::pair { + 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}; +} + +class Lfm2MLP final : public nn::Module { + public: + Lfm2MLP() = default; + Lfm2MLP(const std::string& name, const Lfm2Config& cfg) : nn::Module(name) { + w1_ = reg("w1", cfg.hidden_size, cfg.intermediate_size, false, cfg.linear_impl_type); + w3_ = reg("w3", cfg.hidden_size, cfg.intermediate_size, false, cfg.linear_impl_type); + w2_ = reg("w2", cfg.intermediate_size, cfg.hidden_size, false, cfg.linear_impl_type); + configureLfm2KaiW4A32Threads(w1_); + configureLfm2KaiW4A32Threads(w3_); + configureLfm2KaiW4A32Threads(w2_); + silu_ = reg("silu"); + } + std::vector forward(const std::vector& inputs, const std::vector&) override { + std::vector gate_up; + if (inputs[0].rank() >= 2 && inputs[0].size(-2) == 1) { + const std::array gate_up_ops = {w1_.impl()->getInstancedOp(), w3_.impl()->getInstancedOp()}; + if (cpu::CPULinearOp::tryForwardSharedInputKaiM1(inputs[0], gate_up_ops.data(), gate_up_ops.size(), gate_up)) { + return {w2_(silu_(gate_up[0]) * gate_up[1])}; + } + } + return {w2_(silu_(w1_(inputs[0])) * w3_(inputs[0]))}; + } + + private: + nn::Linear w1_; + nn::Linear w2_; + nn::Linear w3_; + 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; + q_proj_ = reg("q_proj", hidden_size_, query_heads_ * head_dim_, false, cfg.linear_impl_type); + k_proj_ = reg("k_proj", hidden_size_, kv_heads_ * head_dim_, false, cfg.linear_impl_type); + v_proj_ = reg("v_proj", hidden_size_, kv_heads_ * head_dim_, false, cfg.linear_impl_type); + out_proj_ = reg("out_proj", query_heads_ * head_dim_, hidden_size_, false, cfg.linear_impl_type); + configureLfm2KaiW4A32Threads(q_proj_); + configureLfm2KaiW4A32Threads(k_proj_); + configureLfm2KaiW4A32Threads(v_proj_); + configureLfm2KaiW4A32Threads(out_proj_); + 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_); + } + + 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]; + std::vector qkv; + if (sequence == 1) { + const std::array qkv_ops = { + q_proj_.impl()->getInstancedOp(), k_proj_.impl()->getInstancedOp(), v_proj_.impl()->getInstancedOp()}; + (void)cpu::CPULinearOp::tryForwardSharedInputKaiM1(x, qkv_ops.data(), qkv_ops.size(), qkv); + } + if (qkv.empty()) { qkv = {q_proj_(x), k_proj_(x), v_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 = nn::llm_components::groupedQueryAttentionDirectEager(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::Linear q_proj_; + nn::Linear k_proj_; + nn::Linear v_proj_; + nn::Linear out_proj_; + nn::RMSNorm q_layernorm_; + nn::RMSNorm k_layernorm_; + nn::RoPE q_rope_; + nn::RoPE k_rope_; +}; + +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", hidden_size_, 3 * hidden_size_, cfg.conv_bias, cfg.linear_impl_type); + conv_ = reg("conv", hidden_size_, hidden_size_, kernel_size_, 1, 0, 1, hidden_size_, cfg.conv_bias); + out_proj_ = reg("out_proj", hidden_size_, hidden_size_, cfg.conv_bias, cfg.linear_impl_type); + configureLfm2KaiW4A32Threads(in_proj_); + configureLfm2KaiW4A32Threads(out_proj_); + } + + // 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 conv_weight = conv_.weight(); + if (conv_weight.dtype() != kFloat32 || conv_weight.device() != kCPU || !conv_weight.isContiguous()) { + throw std::invalid_argument("LFM2 short convolution requires contiguous float32 CPU weights"); + } + auto convolved = Tensor::empty({batch, sequence, hidden_size_}, kFloat32, kCPU).alloc(); + cpu::gdn::depthwiseCausalConvHistoryFirstF32(bx.ptr(), conv_weight.ptr(), state_.ptr(), + convolved.ptr(), batch, sequence, hidden_size_, kernel_size_); + static const bool trace_activation = [] { + const char* value = std::getenv("MLLM_LFM2_SHORT_CONV_TRACE"); + return value != nullptr && value[0] == '1' && value[1] == '\0'; + }(); + if (trace_activation) { + static std::atomic activated{false}; + if (!activated.exchange(true, std::memory_order_relaxed)) { + std::fprintf(stderr, "MLLM_LFM2_SHORT_CONV_REUSE_ACTIVATED k=%d channels=%d\n", kernel_size_, hidden_size_); + } + } + return {out_proj_(c * convolved)}; + } + + private: + int32_t hidden_size_ = 0; + int32_t kernel_size_ = 0; + nn::Linear in_proj_; + nn::Conv1D 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", cfg.hidden_size, cfg.vocab_size, false, cfg.linear_impl_type); + configureLfm2KaiW4A32Threads(lm_head_); + 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 { + 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..8543197bc --- /dev/null +++ b/mllm/models/lfm2/tokenization_lfm2.hpp @@ -0,0 +1,298 @@ +// 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/tokenizers/AutoTokenizer.hpp" +#include "mllm/preprocessor/tokenizers/BPE.hpp" +#include "mllm/preprocessor/tokenizers/Unicode.hpp" + +namespace mllm::models::lfm2 { + +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); } + + private: + static bool continuation(unsigned char byte) { return byte >= 0x80 && byte <= 0xBF; } + static size_t length(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 validSecond(unsigned char lead, unsigned char second) { + if (!continuation(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) { + static constexpr std::string_view replacement = "\xEF\xBF\xBD"; + std::string output; + size_t offset = 0; + while (offset < pending_.size()) { + const auto lead = static_cast(pending_[offset]); + const auto count = length(lead); + if (count == 1) { + output.push_back(pending_[offset++]); + continue; + } + if (count == 0) { + output.append(replacement); + ++offset; + continue; + } + const auto available = pending_.size() - offset; + bool valid = available < 2 || validSecond(lead, static_cast(pending_[offset + 1])); + for (size_t index = 2; valid && index < std::min(available, count); ++index) { + valid = continuation(static_cast(pending_[offset + index])); + } + if (!valid) { + output.append(replacement); + ++offset; + } else if (available < count) { + if (flush) { + output.append(replacement); + offset = pending_.size(); + } + break; + } else { + output.append(pending_, offset, count); + offset += count; + } + } + pending_.erase(0, offset); + return output; + } + std::string pending_; +}; + +// 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/nn/llm_components/GroupedQueryAttention.hpp b/mllm/nn/llm_components/GroupedQueryAttention.hpp index 91a6bbbe7..378baf2c8 100644 --- a/mllm/nn/llm_components/GroupedQueryAttention.hpp +++ b/mllm/nn/llm_components/GroupedQueryAttention.hpp @@ -3,15 +3,127 @@ #pragma once +#include #include #include +#include #include +#include +#include "mllm/core/Parallel.hpp" #include "mllm/core/Tensor.hpp" #include "mllm/nn/Functional.hpp" 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"); + } +} + +// Direct strided implementation used when a model requires the established +// eager accumulation order. It shares KV heads without materializing an +// expanded cache and parallelizes independent batch/query-head jobs. +inline Tensor groupedQueryAttentionDirectEager(const Tensor& query, const Tensor& key, const Tensor& value) { + validateGroupedQueryAttention(query, key, 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]; + + auto output = Tensor::zeros({q_shape[0], q_shape[1], q_shape[2], v_shape[3]}, value.dtype(), kCPU); + auto compute = [&]() { + const int32_t jobs = q_shape[0] * q_shape[1]; + 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}); + } + } + } + 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; + 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 += static_cast(query_rows[query_row][dim]) * static_cast(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) * static_cast(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] = static_cast(accumulated); + } + } + } + MLLM_AUTO_PARALLEL_FOR_END() + }; + if (query.dtype() == kFloat32) { + compute.template operator()(); + } else { + compute.template operator()(); + } + return output; +} + 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 +176,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/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index 90ce8037b..28e199d91 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -40,6 +40,20 @@ target_include_directories(Mllm-Test-MiniCPM5-Model PRIVATE ${MLLM_INCLUDE_DIR}) target_compile_definitions(Mllm-Test-MiniCPM5-Model PRIVATE MINICPM5_EXAMPLE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../../examples/minicpm5") +add_executable(Mllm-Test-Lfm2-Config Lfm2ConfigTest.cpp) +target_link_libraries(Mllm-Test-Lfm2-Config PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-Lfm2-Config PRIVATE ${MLLM_INCLUDE_DIR}) +target_compile_definitions(Mllm-Test-Lfm2-Config + PRIVATE LFM2_EXAMPLE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../../examples/lfm2") + +add_executable(Mllm-Test-Lfm2-Tokenizer Lfm2TokenizerTest.cpp) +target_link_libraries(Mllm-Test-Lfm2-Tokenizer PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-Lfm2-Tokenizer PRIVATE ${MLLM_INCLUDE_DIR}) + +add_executable(Mllm-Test-Lfm2-ShortConv Lfm2ShortConvTest.cpp) +target_link_libraries(Mllm-Test-Lfm2-ShortConv PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-Lfm2-ShortConv PRIVATE ${MLLM_INCLUDE_DIR}) + add_executable(Mllm-Test-CPUContiguousOp ContiguousOpTest.cpp) target_link_libraries(Mllm-Test-CPUContiguousOp PRIVATE gtest_main MllmRT MllmCPUBackend) target_include_directories(Mllm-Test-CPUContiguousOp PRIVATE ${MLLM_INCLUDE_DIR}) diff --git a/tests/cpu/Lfm2ConfigTest.cpp b/tests/cpu/Lfm2ConfigTest.cpp new file mode 100644 index 000000000..8741eefd4 --- /dev/null +++ b/tests/cpu/Lfm2ConfigTest.cpp @@ -0,0 +1,55 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. +#include + +#include + +#include "mllm/mllm.hpp" +#include "mllm/models/lfm2/configuration_lfm2.hpp" +#include "mllm/models/lfm2/modeling_lfm2.hpp" + +namespace { + +auto loadConfig() -> mllm::models::lfm2::Lfm2Config { + return mllm::models::lfm2::Lfm2Config(std::string(LFM2_EXAMPLE_DIR) + "/config_2.6B_w4a32_kai.json"); +} + +} // namespace + +TEST(Lfm2ConfigTest, Official26BContractUsesCompactAttentionSlots) { + const auto config = loadConfig(); + EXPECT_TRUE(mllm::models::lfm2::matchesOfficialRuntimeContract(config)); + EXPECT_EQ(config.numAttentionLayers(), 8); + EXPECT_EQ(config.numConvLayers(), 22); + EXPECT_EQ(config.attentionSlotForPhysicalLayer(2), 0); + EXPECT_EQ(config.attentionSlotForPhysicalLayer(27), 7); + EXPECT_THROW((void)config.attentionSlotForPhysicalLayer(0), std::invalid_argument); +} + +TEST(Lfm2ConfigTest, NativeKVCacheUsesEightHeadsPerLogicalSlot) { + mllm::initializeContext(); + const auto config = loadConfig(); + auto model = mllm::models::lfm2::Lfm2ForCausalLM(config); + EXPECT_EQ(model.kvCache().getLayerNums(), 8); + EXPECT_EQ(model.kvCache().kvHeads(), 8); + EXPECT_EQ(model.kvCache().headDim(), 64); + EXPECT_EQ(model.kvCache().maxCacheLength(), 2048); + EXPECT_NO_THROW(model.resetState()); + EXPECT_EQ(model.kvCache().getCurrentSeqCnt(0), 0); + mllm::shutdownContext(); +} + +TEST(Lfm2ConfigTest, RejectsPhysicalLayerScheduleDrift) { + auto config = loadConfig(); + config.layer_types[0] = "full_attention"; + EXPECT_FALSE(mllm::models::lfm2::matchesOfficialRuntimeContract(config)); +} + +TEST(Lfm2ConfigTest, KaiW4A32ThreadCapsSeparateDecodeAndPrefill) { + using mllm::cpu::detail::kaiW4A32ThreadCount; + EXPECT_EQ(kaiW4A32ThreadCount(1, 8, 4, 6), 4); + EXPECT_EQ(kaiW4A32ThreadCount(28, 8, 4, 6), 6); + EXPECT_EQ(kaiW4A32ThreadCount(1, 2, 4, 6), 2); + EXPECT_EQ(kaiW4A32ThreadCount(28, 4, 4, 6), 4); + EXPECT_EQ(kaiW4A32ThreadCount(1, 8, 0, 0), 8); +} diff --git a/tests/cpu/Lfm2ShortConvTest.cpp b/tests/cpu/Lfm2ShortConvTest.cpp new file mode 100644 index 000000000..315c5c723 --- /dev/null +++ b/tests/cpu/Lfm2ShortConvTest.cpp @@ -0,0 +1,73 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. +#include + +#include +#include +#include + +#include "mllm/mllm.hpp" +#include "mllm/models/lfm2/modeling_lfm2.hpp" + +namespace { + +class Lfm2ShortConvTest : public testing::Test { + protected: + static void SetUpTestSuite() { mllm::initializeContext(); } +}; + +auto tensor(const std::string& name, const mllm::Tensor::shape_t& shape, const std::vector& values) -> mllm::Tensor { + auto result = mllm::Tensor::empty(shape, mllm::kFloat32, mllm::kCPU).setMemType(mllm::kParamsNormal).setName(name).alloc(); + EXPECT_EQ(result.numel(), values.size()); + std::copy(values.begin(), values.end(), result.ptr()); + return result; +} + +auto input(const std::vector& values) -> mllm::Tensor { + auto result = mllm::Tensor::empty({1, static_cast(values.size()), 1}, mllm::kFloat32, mllm::kCPU).alloc(); + std::copy(values.begin(), values.end(), result.ptr()); + return result; +} + +auto shortConv() -> mllm::models::lfm2::Lfm2ShortConv { + mllm::models::lfm2::Lfm2Config config; + config.hidden_size = 1; + config.conv_L_cache = 3; + config.conv_bias = false; + // Keep this semantics-only test portable. kDefault selects the ARM-only + // MllmBlas fallback for the deliberately tiny K=1 geometry on non-BLAS x86. + config.linear_impl_type = mllm::aops::LinearImplTypes::kGGUF; + auto module = mllm::models::lfm2::Lfm2ShortConv("unit", config); + auto parameters = mllm::ParameterFile::create(); + parameters->push("unit.in_proj.weight", tensor("unit.in_proj.weight", {3, 1}, {1.0F, 1.0F, 1.0F})); + parameters->push("unit.conv.weight", tensor("unit.conv.weight", {1, 1, 3}, {1.0F, 2.0F, 3.0F})); + parameters->push("unit.out_proj.weight", tensor("unit.out_proj.weight", {1, 1}, {1.0F})); + module.load(parameters); + return module; +} + +auto values(mllm::Tensor output) -> std::vector { + output = output.contiguous(); + return {output.ptr(), output.ptr() + output.numel()}; +} + +TEST_F(Lfm2ShortConvTest, ChunkedPrefillAndDecodeMatchOneShotCausalConvolution) { + auto chunked = shortConv(); + auto prefill = values(chunked(input({1.0F, 2.0F}))[0]); + auto decode = values(chunked(input({3.0F}))[0]); + EXPECT_EQ(prefill, (std::vector{3.0F, 28.0F})); + EXPECT_EQ(decode, (std::vector{108.0F})); + EXPECT_EQ(values(chunked.state()), (std::vector{4.0F, 9.0F})); + + auto one_shot = shortConv(); + EXPECT_EQ(values(one_shot(input({1.0F, 2.0F, 3.0F}))[0]), (std::vector{3.0F, 28.0F, 108.0F})); +} + +TEST_F(Lfm2ShortConvTest, ResetClearsTheTwoRequiredHistoricalSamples) { + auto module = shortConv(); + (void)module(input({2.0F})); + module.resetState(1); + EXPECT_EQ(values(module.state()), (std::vector{0.0F, 0.0F})); +} + +} // namespace diff --git a/tests/cpu/Lfm2TokenizerTest.cpp b/tests/cpu/Lfm2TokenizerTest.cpp new file mode 100644 index 000000000..ec7a84b9e --- /dev/null +++ b/tests/cpu/Lfm2TokenizerTest.cpp @@ -0,0 +1,62 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. +#include + +#include +#include +#include + +#include "mllm/mllm.hpp" +#include "mllm/models/lfm2/tokenization_lfm2.hpp" + +TEST(Lfm2TokenizerTest, GroupsDigitsInRunsOfAtMostThree) { + std::vector pieces; + ASSERT_TRUE(mllm::models::lfm2::tokenizerRegex("1234567", pieces)); + EXPECT_EQ(pieces, (std::vector{L"123", L"456", L"7"})); +} + +TEST(Lfm2TokenizerTest, GenerationPromptEndsAtThinkingTokenWithoutNewline) { + const auto text = mllm::models::lfm2::Lfm2Message::render({.prompt = "Hello"}); + EXPECT_EQ(text.substr(text.size() - 7), ""); + EXPECT_EQ(text.find("<|startoftext|>"), 0); + EXPECT_EQ(text.find("<|im_start|>assistant\n"), text.size() - 29); +} + +TEST(Lfm2TokenizerTest, RendersPinnedSystemAndRawToolSchemaContract) { + const auto text = mllm::models::lfm2::Lfm2Message::render( + {.prompt = "Weather?", .system_prompt = "Be concise.", .tools = {R"({"type": "function"})"}}); + EXPECT_EQ(text, "<|startoftext|><|im_start|>system\nBe concise.\nList of tools: [{\"type\": \"function\"}]<|im_end|>\n" + "<|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n"); +} + +TEST(Lfm2TokenizerTest, MatchesPinnedCheckpointOracleWhenProvided) { + const char* tokenizer_path = std::getenv("MLLM_LFM2_TOKENIZER_JSON"); + if (tokenizer_path == nullptr) GTEST_SKIP() << "set MLLM_LFM2_TOKENIZER_JSON to run checkpoint oracle"; + mllm::initializeContext(); + auto tokenizer = mllm::models::lfm2::Lfm2Tokenizer(tokenizer_path); + auto input = tokenizer.convertMessage({.prompt = "Hello"}).at("sequence"); + const std::vector expected = {124894, 124899, 5922, 207, 35808, 124900, 207, 124899, 63514, 207, 124901}; + ASSERT_EQ(input.shape()[1], expected.size()); + for (size_t index = 0; index < expected.size(); ++index) EXPECT_EQ(input.ptr()[index], expected[index]); + + auto tool_input = + tokenizer.convertMessage({.prompt = "Weather?", .system_prompt = "Be concise.", .tools = {R"({"type": "function"})"}}) + .at("sequence"); + const std::vector tool_expected = {124894, 124899, 23630, 207, 4184, 55911, 318, 3120, 302, 5985, + 34, 66155, 5882, 6380, 496, 5545, 66212, 124900, 207, 124899, + 5922, 207, 97056, 39, 124900, 207, 124899, 63514, 207, 124901}; + ASSERT_EQ(tool_input.shape()[1], tool_expected.size()); + for (size_t index = 0; index < tool_expected.size(); ++index) { + EXPECT_EQ(tool_input.ptr()[index], tool_expected[index]); + } + + const std::string multilingual = "你好 LFM2.5!"; + const auto ordinary_tokens = tokenizer.tokenize(multilingual); + const auto ordinary_ids = tokenizer.convert2Ids(ordinary_tokens); + std::string reconstructed; + for (int32_t index = 0; index < ordinary_ids.shape()[1]; ++index) { + reconstructed += tokenizer.detokenizeBytes(ordinary_ids.ptr()[index]); + } + EXPECT_EQ(reconstructed, multilingual); + mllm::shutdownContext(); +} diff --git a/tests/cpu/Qwen35GDNConvTest.cpp b/tests/cpu/Qwen35GDNConvTest.cpp index 46bcaa079..9816e2053 100644 --- a/tests/cpu/Qwen35GDNConvTest.cpp +++ b/tests/cpu/Qwen35GDNConvTest.cpp @@ -23,6 +23,7 @@ namespace { using mllm::cpu::gdn::depthwiseCausalConvF32; +using mllm::cpu::gdn::depthwiseCausalConvHistoryFirstF32; // Deterministic index-derived fill. No RNG, so every host reproduces the same // bytes without carrying a seed through the evidence record. @@ -61,6 +62,29 @@ void referenceDepthwiseCausalConv(const std::vector& input, const std::ve } } +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]; + } + } + } +} + struct ConvCase { int batch; int sequence; @@ -140,6 +164,36 @@ TEST(Qwen35GDNConvTest, MatchesScalarReferenceWithChannelTailAtProductionScale) } } +TEST(Qwen35GDNConvTest, HistoryFirstK3MatchesScalarReferenceBitwiseForLfmWidths) { + constexpr int kKernel = 3; + for (int batch : {1, 2}) { + for (int sequence : {1, 2, 28, 225}) { + for (int channels : {1, 3, 4, 5, 2045, 2048}) { + 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 = makeBuffer(state_count, 19); + + auto kernel_state = initial_state; + std::vector kernel_output(element_count, 0.0F); + 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); + + ASSERT_EQ(kernel_output, reference_output) + << "history-first output mismatch for B=" << batch << " S=" << sequence << " C=" << channels; + ASSERT_EQ(kernel_state, reference_state) + << "history-first state mismatch for B=" << batch << " S=" << sequence << " C=" << channels; + } + } + } +} + TEST(Qwen35GDNConvTest, ChunkedPartitionsMatchOneShot) { struct Partition { int channels; diff --git a/tests/nn/GroupedQueryAttentionTest.cpp b/tests/nn/GroupedQueryAttentionTest.cpp index d57b084bc..0a49fd7e6 100644 --- a/tests/nn/GroupedQueryAttentionTest.cpp +++ b/tests/nn/GroupedQueryAttentionTest.cpp @@ -121,6 +121,16 @@ TEST_F(GroupedQueryAttentionTest, MatchesRepeatedKVReference) { expectNear(actual, expected); } +TEST_F(GroupedQueryAttentionTest, DirectEagerMatchesReference) { + 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::llm_components::groupedQueryAttentionDirectEager(query, key, value); + const auto expected = gqaReference(query, key, value); + + expectNear(actual, expected, 1e-6F); +} + TEST_F(GroupedQueryAttentionTest, SupportsOneKVHeadAndRejectsIllegalGeometry) { auto query = sequential({1, 4, 1, 4}, 0.09F); auto key = sequential({1, 1, 2, 4}, 0.12F); @@ -178,6 +188,21 @@ 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), @@ -196,4 +221,16 @@ TEST_F(GroupedQueryAttentionTest, DecodeOpTraceAndSerializationRoundTrip) { EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kGroupedQueryAttentionDecode); } +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::llm_components::groupedQueryAttention(query, key, value); + const auto expected = gqaReference(query.contiguous(), key.contiguous(), value.contiguous()); + expectNear(actual, expected); +} } // namespace From 524b49dce5b4ba3ece620875bc5aa39841f2a4b9 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Fri, 14 Aug 2026 19:27:12 +0800 Subject: [PATCH 02/21] refactor: register LFM2 hybrid compute paths --- examples/lfm2/README.md | 4 + examples/lfm2/test_validators.py | 7 + examples/lfm2/validate_checkpoint.py | 17 +++ mllm/backends/cpu/CPUBackend.cpp | 6 +- .../cpu/ops/CausalDepthwiseConv1DOp.cpp | 68 +++++++++ .../cpu/ops/CausalDepthwiseConv1DOp.hpp | 24 ++++ .../cpu/ops/GroupedQueryAttentionOp.cpp | 102 +++++++++++++ .../cpu/ops/GroupedQueryAttentionOp.hpp | 24 ++++ mllm/backends/cpu/ops/LinearOp.cpp | 103 +------------ mllm/backends/cpu/ops/LinearOp.hpp | 7 - mllm/backends/cpu/ops/ParallelLinearOp.cpp | 118 +++++++++++++++ mllm/backends/cpu/ops/ParallelLinearOp.hpp | 36 +++++ mllm/compile/ir/GeneratedRTTIKind.hpp | 5 +- mllm/compile/ir/NodeRTTIClassOfImpl.hpp | 13 +- mllm/compile/ir/linalg/Op.cpp | 3 + mllm/compile/ir/linalg/Op.hpp | 6 + mllm/compile/ir/rtti_kind_gen.py | 3 + .../jit/binary/LinalgIRSerialization.cpp | 35 ++++- .../jit/binary/LinalgIRSerialization.hpp | 3 + mllm/compile/jit/interpreter/AopsFromJson.cpp | 67 +++++++++ mllm/compile/jit/interpreter/AopsFromJson.hpp | 3 + mllm/core/OpTypes.hpp | 8 ++ mllm/core/aops/CausalDepthwiseConv1DOp.cpp | 82 +++++++++++ mllm/core/aops/CausalDepthwiseConv1DOp.hpp | 63 ++++++++ mllm/core/aops/GroupedQueryAttentionOp.cpp | 58 ++++++++ mllm/core/aops/GroupedQueryAttentionOp.hpp | 50 +++++++ mllm/core/aops/LinearOp.hpp | 2 + mllm/core/aops/ParallelLinearOp.cpp | 104 ++++++++++++++ mllm/core/aops/ParallelLinearOp.hpp | 47 ++++++ mllm/models/lfm2/modeling_lfm2.hpp | 131 ++++++++--------- mllm/models/lfm2/tokenization_lfm2.hpp | 67 +-------- .../models/minicpm5/tokenization_minicpm5.hpp | 76 +--------- mllm/models/qwen3_5/tokenization_qwen3_5.hpp | 85 +---------- mllm/nn/Functional.cpp | 8 ++ mllm/nn/Functional.hpp | 5 + mllm/nn/Nn.hpp | 3 + mllm/nn/layers/CausalDepthwiseConv1D.cpp | 30 ++++ mllm/nn/layers/CausalDepthwiseConv1D.hpp | 24 ++++ mllm/nn/layers/GroupedQueryAttention.cpp | 14 ++ mllm/nn/layers/GroupedQueryAttention.hpp | 19 +++ mllm/nn/layers/ParallelLinear.cpp | 12 ++ mllm/nn/layers/ParallelLinear.hpp | 21 +++ .../llm_components/GroupedQueryAttention.hpp | 84 ----------- mllm/preprocessor/StreamingUtf8Decoder.hpp | 97 +++++++++++++ tests/cpu/CMakeLists.txt | 4 + tests/cpu/Lfm2ConfigTest.cpp | 4 + tests/cpu/Lfm2RegisteredOpsTest.cpp | 135 ++++++++++++++++++ tests/cpu/Lfm2TokenizerTest.cpp | 7 + tests/nn/GroupedQueryAttentionTest.cpp | 43 +++++- 49 files changed, 1443 insertions(+), 494 deletions(-) create mode 100644 mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.cpp create mode 100644 mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.hpp create mode 100644 mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp create mode 100644 mllm/backends/cpu/ops/GroupedQueryAttentionOp.hpp create mode 100644 mllm/backends/cpu/ops/ParallelLinearOp.cpp create mode 100644 mllm/backends/cpu/ops/ParallelLinearOp.hpp create mode 100644 mllm/core/aops/CausalDepthwiseConv1DOp.cpp create mode 100644 mllm/core/aops/CausalDepthwiseConv1DOp.hpp create mode 100644 mllm/core/aops/GroupedQueryAttentionOp.cpp create mode 100644 mllm/core/aops/GroupedQueryAttentionOp.hpp create mode 100644 mllm/core/aops/ParallelLinearOp.cpp create mode 100644 mllm/core/aops/ParallelLinearOp.hpp create mode 100644 mllm/nn/layers/CausalDepthwiseConv1D.cpp create mode 100644 mllm/nn/layers/CausalDepthwiseConv1D.hpp create mode 100644 mllm/nn/layers/GroupedQueryAttention.cpp create mode 100644 mllm/nn/layers/GroupedQueryAttention.hpp create mode 100644 mllm/nn/layers/ParallelLinear.cpp create mode 100644 mllm/nn/layers/ParallelLinear.hpp create mode 100644 mllm/preprocessor/StreamingUtf8Decoder.hpp create mode 100644 tests/cpu/Lfm2RegisteredOpsTest.cpp diff --git a/examples/lfm2/README.md b/examples/lfm2/README.md index 619ef219e..f72128cee 100644 --- a/examples/lfm2/README.md +++ b/examples/lfm2/README.md @@ -83,6 +83,10 @@ 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 diff --git a/examples/lfm2/test_validators.py b/examples/lfm2/test_validators.py index e06af55e3..815256172 100644 --- a/examples/lfm2/test_validators.py +++ b/examples/lfm2/test_validators.py @@ -20,6 +20,13 @@ def test_official_contract_has_266_tensors(self) -> None: 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) diff --git a/examples/lfm2/validate_checkpoint.py b/examples/lfm2/validate_checkpoint.py index 69e2f3e7d..11aad17e1 100644 --- a/examples/lfm2/validate_checkpoint.py +++ b/examples/lfm2/validate_checkpoint.py @@ -33,6 +33,11 @@ "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", @@ -98,6 +103,16 @@ def validate_config(config: dict) -> None: 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: @@ -178,10 +193,12 @@ def main() -> None: 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) diff --git a/mllm/backends/cpu/CPUBackend.cpp b/mllm/backends/cpu/CPUBackend.cpp index 40a5e6a08..db22e417c 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" @@ -25,10 +26,12 @@ #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 +87,8 @@ CPUBackend::CPUBackend() : Backend(kCPU, createCPUAllocator()) { CPUConv2DOpFactory, CPULayerNorm2DOpFactory, CPUInterpolateOpFactory, CPUPadOpFactory, CPUMaskedScatterOpFactory, CPUArgsortOpFactory, CPUCloneOpFactory, CPUAvgPool1dOpFactory, CPUFlashAttention2SwaSinkOpFactory, CPURadixAttnRelaxOpFactory, CPURadixAttnSwaSinkOpFactory, CPUEqualOpFactory, CPUWhereOpFactory, - CPUGatherOpFactory, CPUGroupedQueryAttentionDecodeOpFactory>(); + CPUGatherOpFactory, CPUGroupedQueryAttentionDecodeOpFactory, CPUCausalDepthwiseConv1DOpFactory, + CPUGroupedQueryAttentionOpFactory, CPUParallelLinearOpFactory>(); } CPUBackend::~CPUBackend() { diff --git a/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.cpp b/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.cpp new file mode 100644 index 000000000..1a7b10df0 --- /dev/null +++ b/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.cpp @@ -0,0 +1,68 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.hpp" + +#include +#include +#include +#include +#include + +#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()); } + + if (options_.accumulation_order == aops::CausalDepthwiseConv1DAccumulationOrder::kHistoryFirst) { + gdn::depthwiseCausalConvHistoryFirstF32(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_LFM2_SHORT_CONV_TRACE"); + return value != nullptr && value[0] == '1' && value[1] == '\0'; + }(); + if (trace_activation) { + static std::atomic activated{false}; + if (!activated.exchange(true, std::memory_order_relaxed)) { + std::fprintf(stderr, "MLLM_LFM2_SHORT_CONV_REUSE_ACTIVATED k=%d channels=%d\n", options_.kernel_size, + options_.channels); + } + } + } else { + gdn::depthwiseCausalConvF32(input.ptr(), weight_.ptr(), updated_state.ptr(), output.ptr(), + input.shape()[0], input.shape()[1], input.shape()[2], options_.kernel_size); + } + + 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/GroupedQueryAttentionOp.cpp b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp new file mode 100644 index 000000000..3705b3218 --- /dev/null +++ b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp @@ -0,0 +1,102 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/ops/GroupedQueryAttentionOp.hpp" + +#include +#include +#include +#include +#include + +#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 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]; + + 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}); + } + } + } + + 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; + 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 += static_cast(query_rows[query_row][dim]) * static_cast(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) * static_cast(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] = static_cast(accumulated); + } + } + } + MLLM_AUTO_PARALLEL_FOR_END() +} + +} // namespace + +CPUGroupedQueryAttentionOp::CPUGroupedQueryAttentionOp(const aops::GroupedQueryAttentionOpOptions& options) + : aops::GroupedQueryAttentionOp(options) {} + +void CPUGroupedQueryAttentionOp::forward(const std::vector& inputs, std::vector& outputs) { + if (options_.implementation != aops::GroupedQueryAttentionImplementation::kDirectStrided) { + throw std::invalid_argument("Unsupported CPU GroupedQueryAttention implementation"); + } + if (inputs[0].dtype() == kFloat32) { + groupedQueryAttentionDirectStrided(inputs[0], inputs[1], inputs[2], outputs[0]); + } else { + groupedQueryAttentionDirectStrided(inputs[0], inputs[1], inputs[2], outputs[0]); + } +} + +} // 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 8309acd75..4105c6273 100644 --- a/mllm/backends/cpu/ops/LinearOp.cpp +++ b/mllm/backends/cpu/ops/LinearOp.cpp @@ -1,13 +1,10 @@ // Copyright (c) MLLM Team. // Licensed under the MIT License. -#include #include #include #include #include -#include -#include #if defined(__linux__) #include @@ -17,7 +14,6 @@ #include "mllm/backends/cpu/kernels/Kernels.hpp" #include "mllm/core/DataTypes.hpp" #include "mllm/core/aops/LinearOp.hpp" -#include "mllm/engine/Context.hpp" namespace mllm::cpu { @@ -76,17 +72,9 @@ void traceKaiW4A32PrefillTile(KaiW4A32Tile tile, int m, int k, int n, int thread CPULinearOp::CPULinearOp(const aops::LinearOpOptions& options) : LinearOp(options) {} -void CPULinearOp::setKaiW4A32ThreadCaps(int decode_thread_cap, int prefill_thread_cap) { - if (decode_thread_cap <= 0 || prefill_thread_cap <= 0) { - throw std::invalid_argument("KAI W4A32 thread caps must be positive"); - } - kai_w4a32_decode_thread_cap_ = decode_thread_cap; - kai_w4a32_prefill_thread_cap_ = prefill_thread_cap; -} - int CPULinearOp::kaiW4A32ThreadCount(int m) const { - return detail::kaiW4A32ThreadCount( - m, options_.getThreads(), kai_w4a32_decode_thread_cap_, kai_w4a32_prefill_thread_cap_); + return detail::kaiW4A32ThreadCount(m, options_.getThreads(), options_.kai_w4a32_decode_thread_cap, + options_.kai_w4a32_prefill_thread_cap); } Tensor CPULinearOp::acquireKaiWorkspace(int32_t workspace_size, int m) { @@ -98,93 +86,6 @@ Tensor CPULinearOp::acquireKaiWorkspace(int32_t workspace_size, int m) { return kai_decode_workspace_; } -bool CPULinearOp::tryForwardSharedInputKaiM1(const Tensor& input, const BaseOp::ptr_t* linear_ops, size_t linear_op_count, - std::vector& outputs) { -#if defined(MLLM_HOST_ARCH_ARM64) || defined(MLLM_HOST_ARCH_ARM) - constexpr size_t kMaximumSharedProjections = 3; - constexpr auto kRequiredImpl = aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32; - using KaiHelper = ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk; - constexpr auto kTile = KaiHelper::Tiles::qai8dxp1x8_qsi4c32p8x8_1x8x32; - - if (Context::instance().thisThread()->trace_mode || input.isNil() || input.device() != kCPU || input.dtype() != kFloat32 - || !input.isContiguous() || input.rank() < 2 || input.size(-2) != 1 || input.size(-1) <= 0 || !outputs.empty() - || linear_ops == nullptr || linear_op_count < 2 || linear_op_count > kMaximumSharedProjections) { - return false; - } - - const auto input_shape = input.shape(); - for (size_t index = 0; index + 2 < input_shape.size(); ++index) { - if (input_shape[index] != 1) { return false; } - } - - const int32_t K = input.size(-1); - int32_t thread_count = 0; - std::array ops{}; - for (size_t index = 0; index < linear_op_count; ++index) { - auto* op = dynamic_cast(linear_ops[index].get()); - if (op == nullptr || op->getDevice() != kCPU || op->options_.impl_type != kRequiredImpl || op->options_.bias - || op->options_.in_channels != K || op->options_.out_channels <= 0 || op->weight_.isNil() - || op->weight_.device() != kCPU || op->options_.getThreads() <= 0) { - return false; - } - if (index == 0) { - thread_count = op->kaiW4A32ThreadCount(1); - } else if (op->kaiW4A32ThreadCount(1) != thread_count) { - return false; - } - ops[index] = op; - } - - std::vector prepared_outputs; - prepared_outputs.reserve(linear_op_count); - for (size_t index = 0; index < linear_op_count; ++index) { - auto output_shape = input_shape; - output_shape.back() = ops[index]->options_.out_channels; - prepared_outputs.emplace_back(Tensor::empty(output_shape, kFloat32, kCPU).alloc()); - } - - std::array projections{}; - for (size_t index = 0; index < linear_op_count; ++index) { - projections[index] = { - .dst = prepared_outputs[index].ptr(), - .packed_weight_bias = reinterpret_cast(ops[index]->weight_.ptr()), - .n = ops[index]->options_.out_channels, - }; - } - - KaiHelper kai_helper; - const size_t workspace_size = kai_helper.workspace_size(1, K, kTile); - if (workspace_size == 0 || workspace_size > static_cast(std::numeric_limits::max())) { return false; } - auto workspace = ops[0]->acquireKaiWorkspace(static_cast(workspace_size), 1); - if (!kai_helper.matmul_shared_input_m1(input.ptr(), projections.data(), linear_op_count, workspace.ptr(), - K, kTile, thread_count)) { - return false; - } - - static const bool trace_activation = [] { - const char* value = std::getenv("MLLM_KAI_SHARED_INPUT_TRACE"); - return value != nullptr && value[0] == '1' && value[1] == '\0'; - }(); - if (trace_activation) { - const uint32_t activation_bit = 1U << static_cast(linear_op_count - 2); - 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 k=%d threads=%d\n", linear_op_count, K, thread_count); - } - } - - outputs = std::move(prepared_outputs); - return true; -#else - (void)input; - (void)linear_ops; - (void)linear_op_count; - (void)outputs; - return false; -#endif -} - void CPULinearOp::load(const ParameterFile::ptr_t& ploader) { switch (ploader->version()) { case ModelFileVersion::kV1: { diff --git a/mllm/backends/cpu/ops/LinearOp.hpp b/mllm/backends/cpu/ops/LinearOp.hpp index 077dc1e7b..315a5dbfa 100644 --- a/mllm/backends/cpu/ops/LinearOp.hpp +++ b/mllm/backends/cpu/ops/LinearOp.hpp @@ -30,19 +30,12 @@ class CPULinearOp final : public aops::LinearOp { void reshape(const std::vector& inputs, std::vector& outputs) override; - static bool tryForwardSharedInputKaiM1(const Tensor& input, const BaseOp::ptr_t* linear_ops, size_t linear_op_count, - std::vector& outputs); - - void setKaiW4A32ThreadCaps(int decode_thread_cap, int prefill_thread_cap); - private: Tensor acquireKaiWorkspace(int32_t workspace_size, int m); [[nodiscard]] int kaiW4A32ThreadCount(int m) const; Tensor kai_decode_workspace_; - int kai_w4a32_decode_thread_cap_ = 0; - int kai_w4a32_prefill_thread_cap_ = 0; }; class CPULinearOpFactory : public TypedOpFactory { diff --git a/mllm/backends/cpu/ops/ParallelLinearOp.cpp b/mllm/backends/cpu/ops/ParallelLinearOp.cpp new file mode 100644 index 000000000..d8dbfed28 --- /dev/null +++ b/mllm/backends/cpu/ops/ParallelLinearOp.cpp @@ -0,0 +1,118 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/ops/ParallelLinearOp.hpp" + +#include +#include +#include +#include +#include +#include + +#include "mllm/backends/cpu/kernels/Kernels.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) { + 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::tryForwardSharedInputKaiM1(const Tensor& input, std::vector& outputs) { +#if defined(MLLM_HOST_ARCH_ARM64) || defined(MLLM_HOST_ARCH_ARM) + constexpr size_t kMaximumSharedProjections = 3; + constexpr auto kRequiredImpl = aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32; + using KaiHelper = ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk; + constexpr auto kTile = KaiHelper::Tiles::qai8dxp1x8_qsi4c32p8x8_1x8x32; + + if (input.isNil() || input.device() != kCPU || input.dtype() != kFloat32 || !input.isContiguous() || input.rank() < 2 + || input.size(-2) != 1 || 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; } + } + for (size_t index = 0; index < weights_.size(); ++index) { + if (weights_[index].isNil() || weights_[index].device() != kCPU || outputs[index].dtype() != kFloat32 + || outputs[index].device() != kCPU) { + return false; + } + } + + const int32_t thread_count = detail::kaiW4A32ThreadCount(1, options_.getThreads(), options_.kai_w4a32_decode_thread_cap, + options_.kai_w4a32_prefill_thread_cap); + 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], + }; + } + + KaiHelper kai_helper; + const size_t workspace_size = kai_helper.workspace_size(1, options_.in_channels, kTile); + if (workspace_size == 0 || workspace_size > static_cast(std::numeric_limits::max())) { return false; } + auto workspace = acquireKaiWorkspace(static_cast(workspace_size)); + if (!kai_helper.matmul_shared_input_m1(input.ptr(), projections.data(), weights_.size(), workspace.ptr(), + options_.in_channels, kTile, thread_count)) { + return false; + } + + static const bool trace_activation = [] { + const char* value = std::getenv("MLLM_KAI_SHARED_INPUT_TRACE"); + return value != nullptr && value[0] == '1' && value[1] == '\0'; + }(); + if (trace_activation) { + const uint32_t activation_bit = 1U << static_cast(weights_.size() - 2); + 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 k=%d threads=%d\n", weights_.size(), options_.in_channels, + thread_count); + } + } + return true; +#else + (void)input; + (void)outputs; + return false; +#endif +} + +void CPUParallelLinearOp::forward(const std::vector& inputs, std::vector& outputs) { + const auto& input = inputs[0]; + if (tryForwardSharedInputKaiM1(input, outputs)) { return; } + for (size_t index = 0; index < fallback_ops_.size(); ++index) { + 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..14a6faba6 --- /dev/null +++ b/mllm/backends/cpu/ops/ParallelLinearOp.hpp @@ -0,0 +1,36 @@ +// 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 tryForwardSharedInputKaiM1(const Tensor& input, std::vector& outputs); + Tensor acquireKaiWorkspace(int32_t workspace_size); + + std::vector> fallback_ops_; + 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..50e9b21cd 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-14 18:50:28 // do not modify this file #pragma once @@ -42,6 +42,9 @@ enum NodeKind : uint32_t { 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..fd59ffe7f 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-14 18:50:28 // do not modify this file #pragma once namespace mllm::ir { @@ -97,6 +97,17 @@ struct NodeRTTIClassOfImpl { 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..150d71c0b 100644 --- a/mllm/compile/ir/linalg/Op.cpp +++ b/mllm/compile/ir/linalg/Op.cpp @@ -67,6 +67,9 @@ 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..09d6831d3 100644 --- a/mllm/compile/ir/linalg/Op.hpp +++ b/mllm/compile/ir/linalg/Op.hpp @@ -36,6 +36,9 @@ class ViewOp; class SplitOp; class FlashAttention2Op; class GroupedQueryAttentionDecodeOp; +class CausalDepthwiseConv1DOp; +class GroupedQueryAttentionOp; +class ParallelLinearOp; class RepeatOp; class PermuteOp; class Conv1DOp; @@ -199,6 +202,9 @@ 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..e7b5c0ace 100644 --- a/mllm/compile/ir/rtti_kind_gen.py +++ b/mllm/compile/ir/rtti_kind_gen.py @@ -248,6 +248,9 @@ def define_lianlg_ir(ir: dict): 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..f3197fe81 100644 --- a/mllm/compile/jit/binary/LinalgIRSerialization.cpp +++ b/mllm/compile/jit/binary/LinalgIRSerialization.cpp @@ -12,6 +12,9 @@ #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" @@ -71,6 +74,9 @@ nlohmann::json dumpLinalgIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { CASE(STFT) CASE(FlashAttention2) CASE(GroupedQueryAttentionDecode) + CASE(CausalDepthwiseConv1D) + CASE(GroupedQueryAttention) + CASE(ParallelLinear) CASE(Repeat) CASE(Permute) CASE(Conv1D) @@ -137,7 +143,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 {}; } diff --git a/mllm/compile/jit/binary/LinalgIRSerialization.hpp b/mllm/compile/jit/binary/LinalgIRSerialization.hpp index 7c6a314ab..705b6f3e3 100644 --- a/mllm/compile/jit/binary/LinalgIRSerialization.hpp +++ b/mllm/compile/jit/binary/LinalgIRSerialization.hpp @@ -38,6 +38,9 @@ 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..449b2c542 100644 --- a/mllm/compile/jit/interpreter/AopsFromJson.cpp +++ b/mllm/compile/jit/interpreter/AopsFromJson.cpp @@ -23,6 +23,9 @@ #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 +109,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 +241,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; @@ -607,6 +622,58 @@ BaseOp::ptr_t __groupedQueryAttentionDecodeFromJson(const nlohmann::json& json) return Context::instance().getBackend(backend)->createOp(OpTypes::kGroupedQueryAttentionDecode, 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) { aops::RepeatOpOptions options; 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..fb88f67c1 100644 --- a/mllm/core/OpTypes.hpp +++ b/mllm/core/OpTypes.hpp @@ -100,6 +100,11 @@ enum class OpTypes : int32_t { // Phase-specific native KV-head attention. kGroupedQueryAttentionDecode = 76, + // 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, @@ -185,6 +190,9 @@ inline std::string optype2Str(OpTypes type) { 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/GroupedQueryAttentionOp.cpp b/mllm/core/aops/GroupedQueryAttentionOp.cpp new file mode 100644 index 000000000..fd0df3631 --- /dev/null +++ b/mllm/core/aops/GroupedQueryAttentionOp.cpp @@ -0,0 +1,58 @@ +// 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"); + } + 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..e6aa99d95 --- /dev/null +++ b/mllm/core/aops/GroupedQueryAttentionOp.hpp @@ -0,0 +1,50 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include "mllm/core/BaseOp.hpp" + +namespace mllm::aops { + +enum class GroupedQueryAttentionImplementation : int32_t { + kDirectStrided = 0, +}; + +inline const char* groupedQueryAttentionImplementation2Str(GroupedQueryAttentionImplementation implementation) { + switch (implementation) { + case GroupedQueryAttentionImplementation::kDirectStrided: return "DirectStrided"; + } + return "Unknown"; +} + +inline GroupedQueryAttentionImplementation str2GroupedQueryAttentionImplementation(const std::string& value) { + if (value == "DirectStrided") return GroupedQueryAttentionImplementation::kDirectStrided; + 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..26142cd9c --- /dev/null +++ b/mllm/core/aops/ParallelLinearOp.cpp @@ -0,0 +1,104 @@ +// 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) {} + +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::load(const ParameterFile::ptr_t& ploader) { + 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"); + } + 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 || options_.out_channels.size() < 2 + || options_.projection_names.size() != options_.out_channels.size()) { + throw std::invalid_argument("ParallelLinear options are incomplete"); + } + 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..bcadf8603 --- /dev/null +++ b/mllm/core/aops/ParallelLinearOp.hpp @@ -0,0 +1,47 @@ +// 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; + + inline const ParallelLinearOpOptions& options() const { return options_; } + + protected: + [[nodiscard]] std::string projectionParameterName(size_t index, const char* suffix) const; + + std::vector weights_; + std::vector biases_; + ParallelLinearOpOptions options_; +}; + +} // namespace mllm::aops diff --git a/mllm/models/lfm2/modeling_lfm2.hpp b/mllm/models/lfm2/modeling_lfm2.hpp index b8fabb0d6..5943026bd 100644 --- a/mllm/models/lfm2/modeling_lfm2.hpp +++ b/mllm/models/lfm2/modeling_lfm2.hpp @@ -2,36 +2,44 @@ // Licensed under the MIT License. #pragma once -#include -#include #include -#include -#include #include #include #include #include -#include "mllm/backends/cpu/ops/LinearOp.hpp" -#include "mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp" #include "mllm/core/Tensor.hpp" #include "mllm/models/ARGeneration.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/llm_components/GroupedQueryAttention.hpp" #include "mllm/nn/lmcache/KVHeadStaticCache.hpp" namespace mllm::models::lfm2 { -inline void configureLfm2KaiW4A32Threads(nn::Linear& linear) { - auto op = std::dynamic_pointer_cast(linear.impl()->getInstancedOp()); - if (op != nullptr) { - // OnePlus 13T source-bound screening keeps I8MM prefill above 80 tok/s - // with six workers, while four workers avoid decode GEMV oversubscription. - op->setKaiW4A32ThreadCaps(4, 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, + // OnePlus 13T source-bound screening keeps I8MM prefill above 80 tok/s + // with six workers, while four workers avoid decode GEMV oversubscription. + .kai_w4a32_decode_thread_cap = 4, + .kai_w4a32_prefill_thread_cap = 6}; +} + +inline auto makeLfm2ParallelLinearOptions(int32_t in_channels, std::vector out_channels, + std::vector projection_names, bool bias, aops::LinearImplTypes impl_type) + -> aops::ParallelLinearOpOptions { + return {.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 = 4, + .kai_w4a32_prefill_thread_cap = 6}; } // Model-level orchestration: materialize the immutable analytical RoPE table. @@ -73,29 +81,20 @@ class Lfm2MLP final : public nn::Module { public: Lfm2MLP() = default; Lfm2MLP(const std::string& name, const Lfm2Config& cfg) : nn::Module(name) { - w1_ = reg("w1", cfg.hidden_size, cfg.intermediate_size, false, cfg.linear_impl_type); - w3_ = reg("w3", cfg.hidden_size, cfg.intermediate_size, false, cfg.linear_impl_type); - w2_ = reg("w2", cfg.intermediate_size, cfg.hidden_size, false, cfg.linear_impl_type); - configureLfm2KaiW4A32Threads(w1_); - configureLfm2KaiW4A32Threads(w3_); - configureLfm2KaiW4A32Threads(w2_); + gate_up_proj_ = reg( + "gate_up_proj", makeLfm2ParallelLinearOptions(cfg.hidden_size, {cfg.intermediate_size, cfg.intermediate_size}, + {"w1", "w3"}, false, cfg.linear_impl_type)); + 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 { - std::vector gate_up; - if (inputs[0].rank() >= 2 && inputs[0].size(-2) == 1) { - const std::array gate_up_ops = {w1_.impl()->getInstancedOp(), w3_.impl()->getInstancedOp()}; - if (cpu::CPULinearOp::tryForwardSharedInputKaiM1(inputs[0], gate_up_ops.data(), gate_up_ops.size(), gate_up)) { - return {w2_(silu_(gate_up[0]) * gate_up[1])}; - } - } - return {w2_(silu_(w1_(inputs[0])) * w3_(inputs[0]))}; + auto gate_up = gate_up_proj_(inputs[0]); + return {w2_(silu_(gate_up[0]) * gate_up[1])}; } private: - nn::Linear w1_; + nn::ParallelLinear gate_up_proj_; nn::Linear w2_; - nn::Linear w3_; nn::SiLU silu_; }; @@ -107,31 +106,26 @@ class Lfm2Attention final : public nn::Module { head_dim_ = cfg.head_dim; query_heads_ = cfg.num_attention_heads; kv_heads_ = cfg.num_key_value_heads; - q_proj_ = reg("q_proj", hidden_size_, query_heads_ * head_dim_, false, cfg.linear_impl_type); - k_proj_ = reg("k_proj", hidden_size_, kv_heads_ * head_dim_, false, cfg.linear_impl_type); - v_proj_ = reg("v_proj", hidden_size_, kv_heads_ * head_dim_, false, cfg.linear_impl_type); - out_proj_ = reg("out_proj", query_heads_ * head_dim_, hidden_size_, false, cfg.linear_impl_type); - configureLfm2KaiW4A32Threads(q_proj_); - configureLfm2KaiW4A32Threads(k_proj_); - configureLfm2KaiW4A32Threads(v_proj_); - configureLfm2KaiW4A32Threads(out_proj_); + qkv_proj_ = reg( + "qkv_proj", + makeLfm2ParallelLinearOptions(hidden_size_, {query_heads_ * head_dim_, kv_heads_ * head_dim_, kv_heads_ * head_dim_}, + {"q_proj", "k_proj", "v_proj"}, false, cfg.linear_impl_type)); + 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::GroupedQueryAttentionOpOptions{.implementation = 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]; - std::vector qkv; - if (sequence == 1) { - const std::array qkv_ops = { - q_proj_.impl()->getInstancedOp(), k_proj_.impl()->getInstancedOp(), v_proj_.impl()->getInstancedOp()}; - (void)cpu::CPULinearOp::tryForwardSharedInputKaiM1(x, qkv_ops.data(), qkv_ops.size(), qkv); - } - if (qkv.empty()) { qkv = {q_proj_(x), k_proj_(x), v_proj_(x)}; } + 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_}); @@ -147,7 +141,7 @@ class Lfm2Attention final : public nn::Module { // 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 = nn::llm_components::groupedQueryAttentionDirectEager(query, updated[0], updated[1]); + 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)}; } @@ -159,14 +153,13 @@ class Lfm2Attention final : public nn::Module { int32_t head_dim_ = 0; int32_t query_heads_ = 0; int32_t kv_heads_ = 0; - nn::Linear q_proj_; - nn::Linear k_proj_; - nn::Linear v_proj_; + 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 { @@ -175,11 +168,12 @@ class Lfm2ShortConv final : public nn::Module { 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", hidden_size_, 3 * hidden_size_, cfg.conv_bias, cfg.linear_impl_type); - conv_ = reg("conv", hidden_size_, hidden_size_, kernel_size_, 1, 0, 1, hidden_size_, cfg.conv_bias); - out_proj_ = reg("out_proj", hidden_size_, hidden_size_, cfg.conv_bias, cfg.linear_impl_type); - configureLfm2KaiW4A32Threads(in_proj_); - configureLfm2KaiW4A32Threads(out_proj_); + 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 @@ -202,23 +196,8 @@ class Lfm2ShortConv final : public nn::Module { 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 conv_weight = conv_.weight(); - if (conv_weight.dtype() != kFloat32 || conv_weight.device() != kCPU || !conv_weight.isContiguous()) { - throw std::invalid_argument("LFM2 short convolution requires contiguous float32 CPU weights"); - } - auto convolved = Tensor::empty({batch, sequence, hidden_size_}, kFloat32, kCPU).alloc(); - cpu::gdn::depthwiseCausalConvHistoryFirstF32(bx.ptr(), conv_weight.ptr(), state_.ptr(), - convolved.ptr(), batch, sequence, hidden_size_, kernel_size_); - static const bool trace_activation = [] { - const char* value = std::getenv("MLLM_LFM2_SHORT_CONV_TRACE"); - return value != nullptr && value[0] == '1' && value[1] == '\0'; - }(); - if (trace_activation) { - static std::atomic activated{false}; - if (!activated.exchange(true, std::memory_order_relaxed)) { - std::fprintf(stderr, "MLLM_LFM2_SHORT_CONV_REUSE_ACTIVATED k=%d channels=%d\n", kernel_size_, hidden_size_); - } - } + auto [convolved, updated_state] = conv_(bx, state_); + state_ = std::move(updated_state); return {out_proj_(c * convolved)}; } @@ -226,7 +205,7 @@ class Lfm2ShortConv final : public nn::Module { int32_t hidden_size_ = 0; int32_t kernel_size_ = 0; nn::Linear in_proj_; - nn::Conv1D conv_; + nn::CausalDepthwiseConv1D conv_; nn::Linear out_proj_; Tensor state_; }; @@ -333,8 +312,8 @@ class Lfm2ForCausalLM final : public ARGeneration, public nn::Module { 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", cfg.hidden_size, cfg.vocab_size, false, cfg.linear_impl_type); - configureLfm2KaiW4A32Threads(lm_head_); + 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)); } @@ -361,6 +340,10 @@ class Lfm2ForCausalLM final : public ARGeneration, public nn::Module { *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; } diff --git a/mllm/models/lfm2/tokenization_lfm2.hpp b/mllm/models/lfm2/tokenization_lfm2.hpp index 8543197bc..2a8f52090 100644 --- a/mllm/models/lfm2/tokenization_lfm2.hpp +++ b/mllm/models/lfm2/tokenization_lfm2.hpp @@ -10,77 +10,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::lfm2 { -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); } - - private: - static bool continuation(unsigned char byte) { return byte >= 0x80 && byte <= 0xBF; } - static size_t length(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 validSecond(unsigned char lead, unsigned char second) { - if (!continuation(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) { - static constexpr std::string_view replacement = "\xEF\xBF\xBD"; - std::string output; - size_t offset = 0; - while (offset < pending_.size()) { - const auto lead = static_cast(pending_[offset]); - const auto count = length(lead); - if (count == 1) { - output.push_back(pending_[offset++]); - continue; - } - if (count == 0) { - output.append(replacement); - ++offset; - continue; - } - const auto available = pending_.size() - offset; - bool valid = available < 2 || validSecond(lead, static_cast(pending_[offset + 1])); - for (size_t index = 2; valid && index < std::min(available, count); ++index) { - valid = continuation(static_cast(pending_[offset + index])); - } - if (!valid) { - output.append(replacement); - ++offset; - } else if (available < count) { - if (flush) { - output.append(replacement); - offset = pending_.size(); - } - break; - } else { - output.append(pending_, offset, count); - offset += count; - } - } - pending_.erase(0, offset); - return output; - } - std::string pending_; -}; +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. 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..db70e6d57 100644 --- a/mllm/nn/Functional.cpp +++ b/mllm/nn/Functional.cpp @@ -7,6 +7,7 @@ #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" @@ -93,6 +94,13 @@ Tensor groupedQueryAttentionDecode(const Tensor& query, const Tensor& key, const aops::GroupedQueryAttentionDecodeOpOptions{}, {query, key, value})[0]; } +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) { return Context::instance().buildOpAndSubmitTask(OpTypes::kSoftmax, aops::SoftmaxOpOptions{.axis = dim}, {x})[0]; } 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/Nn.hpp b/mllm/nn/Nn.hpp index f47481c01..49c653679 100644 --- a/mllm/nn/Nn.hpp +++ b/mllm/nn/Nn.hpp @@ -13,6 +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/CausalDepthwiseConv1D.hpp" // IWYU pragma: export +#include "mllm/nn/layers/GroupedQueryAttention.hpp" // IWYU pragma: export #include "mllm/nn/layers/GroupedQueryAttentionDecode.hpp" // IWYU pragma: export #include "mllm/nn/layers/QuickGELU.hpp" // IWYU pragma: export #include "mllm/nn/layers/ReLU.hpp" // IWYU pragma: export @@ -25,6 +27,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..6039c5867 --- /dev/null +++ b/mllm/nn/layers/CausalDepthwiseConv1D.cpp @@ -0,0 +1,30 @@ +// 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(const aops::CausalDepthwiseConv1DOpOptions& options) + : Layer(OpTypes::kCausalDepthwiseConv1D, options) {} + +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..f6c79fff1 --- /dev/null +++ b/mllm/nn/layers/CausalDepthwiseConv1D.hpp @@ -0,0 +1,24 @@ +// 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(); + explicit CausalDepthwiseConv1D(const aops::CausalDepthwiseConv1DOpOptions& options); + 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..bce4477a3 --- /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(const aops::GroupedQueryAttentionOpOptions& options) + : Layer(OpTypes::kGroupedQueryAttention, options) {} + +} // namespace mllm::nn diff --git a/mllm/nn/layers/GroupedQueryAttention.hpp b/mllm/nn/layers/GroupedQueryAttention.hpp new file mode 100644 index 000000000..67f00bc37 --- /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(const aops::GroupedQueryAttentionOpOptions& 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..36c739f27 --- /dev/null +++ b/mllm/nn/layers/ParallelLinear.cpp @@ -0,0 +1,12 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/nn/layers/ParallelLinear.hpp" + +namespace mllm::nn { + +ParallelLinear::ParallelLinear() : Layer(OpTypes::kParallelLinear, aops::ParallelLinearOpOptions{}) {} + +ParallelLinear::ParallelLinear(const aops::ParallelLinearOpOptions& options) : Layer(OpTypes::kParallelLinear, options) {} + +} // namespace mllm::nn diff --git a/mllm/nn/layers/ParallelLinear.hpp b/mllm/nn/layers/ParallelLinear.hpp new file mode 100644 index 000000000..11444080d --- /dev/null +++ b/mllm/nn/layers/ParallelLinear.hpp @@ -0,0 +1,21 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include + +#include "mllm/core/aops/ParallelLinearOp.hpp" +#include "mllm/nn/Layer.hpp" + +namespace mllm::nn { + +class ParallelLinear : public Layer { + public: + ParallelLinear(); + explicit ParallelLinear(const aops::ParallelLinearOpOptions& options); + + std::vector operator()(const Tensor& input) { return __main({input}); } +}; + +} // namespace mllm::nn diff --git a/mllm/nn/llm_components/GroupedQueryAttention.hpp b/mllm/nn/llm_components/GroupedQueryAttention.hpp index 378baf2c8..35327d029 100644 --- a/mllm/nn/llm_components/GroupedQueryAttention.hpp +++ b/mllm/nn/llm_components/GroupedQueryAttention.hpp @@ -3,14 +3,10 @@ #pragma once -#include #include #include -#include #include -#include -#include "mllm/core/Parallel.hpp" #include "mllm/core/Tensor.hpp" #include "mllm/nn/Functional.hpp" @@ -44,86 +40,6 @@ inline void validateGroupedQueryAttention(const Tensor& query, const Tensor& key } } -// Direct strided implementation used when a model requires the established -// eager accumulation order. It shares KV heads without materializing an -// expanded cache and parallelizes independent batch/query-head jobs. -inline Tensor groupedQueryAttentionDirectEager(const Tensor& query, const Tensor& key, const Tensor& value) { - validateGroupedQueryAttention(query, key, 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]; - - auto output = Tensor::zeros({q_shape[0], q_shape[1], q_shape[2], v_shape[3]}, value.dtype(), kCPU); - auto compute = [&]() { - const int32_t jobs = q_shape[0] * q_shape[1]; - 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}); - } - } - } - 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; - 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 += static_cast(query_rows[query_row][dim]) * static_cast(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) * static_cast(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] = static_cast(accumulated); - } - } - } - MLLM_AUTO_PARALLEL_FOR_END() - }; - if (query.dtype() == kFloat32) { - compute.template operator()(); - } else { - compute.template operator()(); - } - return output; -} - inline Tensor groupedQueryAttentionEager(const Tensor& query, const Tensor& key, const Tensor& value) { const auto q_shape = query.shape(); const auto k_shape = key.shape(); 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/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index 28e199d91..aae1210df 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -54,6 +54,10 @@ add_executable(Mllm-Test-Lfm2-ShortConv Lfm2ShortConvTest.cpp) target_link_libraries(Mllm-Test-Lfm2-ShortConv PRIVATE gtest_main MllmCPUBackend) target_include_directories(Mllm-Test-Lfm2-ShortConv PRIVATE ${MLLM_INCLUDE_DIR}) +add_executable(Mllm-Test-Lfm2-RegisteredOps Lfm2RegisteredOpsTest.cpp) +target_link_libraries(Mllm-Test-Lfm2-RegisteredOps PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-Lfm2-RegisteredOps PRIVATE ${MLLM_INCLUDE_DIR}) + add_executable(Mllm-Test-CPUContiguousOp ContiguousOpTest.cpp) target_link_libraries(Mllm-Test-CPUContiguousOp PRIVATE gtest_main MllmRT MllmCPUBackend) target_include_directories(Mllm-Test-CPUContiguousOp PRIVATE ${MLLM_INCLUDE_DIR}) diff --git a/tests/cpu/Lfm2ConfigTest.cpp b/tests/cpu/Lfm2ConfigTest.cpp index 8741eefd4..f7faff726 100644 --- a/tests/cpu/Lfm2ConfigTest.cpp +++ b/tests/cpu/Lfm2ConfigTest.cpp @@ -4,6 +4,7 @@ #include +#include "mllm/backends/cpu/ops/LinearOp.hpp" #include "mllm/mllm.hpp" #include "mllm/models/lfm2/configuration_lfm2.hpp" #include "mllm/models/lfm2/modeling_lfm2.hpp" @@ -36,6 +37,9 @@ TEST(Lfm2ConfigTest, NativeKVCacheUsesEightHeadsPerLogicalSlot) { EXPECT_EQ(model.kvCache().maxCacheLength(), 2048); EXPECT_NO_THROW(model.resetState()); EXPECT_EQ(model.kvCache().getCurrentSeqCnt(0), 0); + model.kvCache().setCurrentSeqCnt(1); + auto sequence = mllm::Tensor::zeros({1, 1}, mllm::kInt64, mllm::kCPU); + EXPECT_THROW((void)model.forward({{"sequence", sequence}}, {}), std::invalid_argument); mllm::shutdownContext(); } diff --git a/tests/cpu/Lfm2RegisteredOpsTest.cpp b/tests/cpu/Lfm2RegisteredOpsTest.cpp new file mode 100644 index 000000000..480d6e09c --- /dev/null +++ b/tests/cpu/Lfm2RegisteredOpsTest.cpp @@ -0,0 +1,135 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#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 Lfm2RegisteredOpsTest : public testing::Test { + protected: + static void SetUpTestSuite() { mllm::initializeContext(); } +}; + +class CausalConvTraceModule final : public mllm::nn::Module { + public: + CausalConvTraceModule() : Module("causal_conv_trace") { + conv_ = reg("conv", 4, 3, false, true, + 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}; + } + + private: + mllm::nn::CausalDepthwiseConv1D conv_; +}; + +class ParallelLinearModule final : public mllm::nn::Module { + public: + explicit ParallelLinearModule(std::string name) : Module(std::move(name)) { + projections_ = reg( + "pair", mllm::aops::ParallelLinearOpOptions{.in_channels = 2, + .out_channels = {2, 1}, + .projection_names = {"left", "right"}, + .bias = false, + .impl_type = mllm::aops::LinearImplTypes::kGGUF, + .kai_w4a32_decode_thread_cap = 4, + .kai_w4a32_prefill_thread_cap = 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 tensor = Tensor::empty(shape, mllm::kFloat32, mllm::kCPU).setMemType(mllm::kParamsNormal).setName(name).alloc(); + std::copy(values.begin(), values.end(), tensor.ptr()); + return tensor; +} + +TEST_F(Lfm2RegisteredOpsTest, CausalConvTracesAndSerializesStateSemantics) { + CausalConvTraceModule module; + auto ir_context = mllm::ir::trace(module, Tensor::empty({1, 2, 4}, mllm::kFloat32, mllm::kCPU), + Tensor::empty({1, 4, 2}, mllm::kFloat32, mllm::kCPU)); + auto op = findOp(ir_context->topLevelOp()); + ASSERT_NE(op, nullptr); + const auto options = mllm::jit::binary::dumpLinalgIROptions(op); + EXPECT_EQ(options.at("channels"), 4); + EXPECT_EQ(options.at("kernel_size"), 3); + 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); + EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kCausalDepthwiseConv1D); +} + +TEST_F(Lfm2RegisteredOpsTest, ParallelLinearOwnsSiblingParametersAndFallsBackCorrectly) { + ParallelLinearModule module("parallel_eager"); + auto parameters = mllm::ParameterFile::create(); + parameters->push("parallel_eager.left.weight", parameter("parallel_eager.left.weight", {2, 2}, {1.0F, 2.0F, 3.0F, 4.0F})); + parameters->push("parallel_eager.right.weight", parameter("parallel_eager.right.weight", {1, 2}, {5.0F, 6.0F})); + module.load(parameters); + + auto input = Tensor::empty({1, 1, 2}, mllm::kFloat32, mllm::kCPU).alloc(); + input.ptr()[0] = 2.0F; + input.ptr()[1] = 3.0F; + const auto outputs = module(input); + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].shape(), (Tensor::shape_t{1, 1, 2})); + EXPECT_FLOAT_EQ(outputs[0].ptr()[0], 8.0F); + EXPECT_FLOAT_EQ(outputs[0].ptr()[1], 18.0F); + EXPECT_FLOAT_EQ(outputs[1].ptr()[0], 28.0F); +} + +TEST_F(Lfm2RegisteredOpsTest, ParallelLinearTracesAndSerializesProjectionContract) { + ParallelLinearModule module("parallel_trace"); + 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); + const auto options = mllm::jit::binary::dumpLinalgIROptions(op); + EXPECT_EQ(options.at("out_channels"), (std::vector{2, 1})); + EXPECT_EQ(options.at("projection_names"), (std::vector{"left", "right"})); + EXPECT_EQ(options.at("kai_w4a32_decode_thread_cap"), 4); + EXPECT_EQ(options.at("kai_w4a32_prefill_thread_cap"), 6); + + const auto restored = mllm::jit::interpreter::aopsFromJson( + nlohmann::json{{"op_type", "ParallelLinear"}, {"backend", "CPU"}, {"op_options", options}}); + ASSERT_NE(restored, nullptr); + EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kParallelLinear); +} + +} // namespace diff --git a/tests/cpu/Lfm2TokenizerTest.cpp b/tests/cpu/Lfm2TokenizerTest.cpp index ec7a84b9e..38df63bb1 100644 --- a/tests/cpu/Lfm2TokenizerTest.cpp +++ b/tests/cpu/Lfm2TokenizerTest.cpp @@ -29,6 +29,13 @@ TEST(Lfm2TokenizerTest, RendersPinnedSystemAndRawToolSchemaContract) { "<|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n"); } +TEST(Lfm2TokenizerTest, StreamsUtf8AcrossTokenBoundaries) { + mllm::models::lfm2::StreamingUtf8Decoder decoder; + EXPECT_EQ(decoder.append("\xF0\x9F"), ""); + EXPECT_EQ(decoder.append("\x98\x80"), "\xF0\x9F\x98\x80"); + EXPECT_EQ(decoder.finish(), ""); +} + TEST(Lfm2TokenizerTest, MatchesPinnedCheckpointOracleWhenProvided) { const char* tokenizer_path = std::getenv("MLLM_LFM2_TOKENIZER_JSON"); if (tokenizer_path == nullptr) GTEST_SKIP() << "set MLLM_LFM2_TOKENIZER_JSON to run checkpoint oracle"; diff --git a/tests/nn/GroupedQueryAttentionTest.cpp b/tests/nn/GroupedQueryAttentionTest.cpp index 0a49fd7e6..43ce947c0 100644 --- a/tests/nn/GroupedQueryAttentionTest.cpp +++ b/tests/nn/GroupedQueryAttentionTest.cpp @@ -37,6 +37,15 @@ class GroupedQueryAttentionDecodeTraceModule final : public mllm::nn::Module { } }; +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::GroupedQueryAttentionDecodeOp::ptr_t findGroupedQueryAttentionDecodeOp(const mllm::ir::node_ptr_t& node) { if (node->isa_()) { return node->cast_(); @@ -50,6 +59,19 @@ mllm::ir::linalg::GroupedQueryAttentionDecodeOp::ptr_t findGroupedQueryAttention return nullptr; } +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 = findGroupedQueryAttentionOp(op)) { return found; } + } + } + return nullptr; +} + Tensor sequential(const Tensor::shape_t& shape, float scale) { auto tensor = Tensor::empty(shape, mllm::kFloat32, mllm::kCPU).alloc(); for (int index = 0; index < tensor.numel(); ++index) { @@ -121,16 +143,33 @@ TEST_F(GroupedQueryAttentionTest, MatchesRepeatedKVReference) { expectNear(actual, expected); } -TEST_F(GroupedQueryAttentionTest, DirectEagerMatchesReference) { +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::llm_components::groupedQueryAttentionDirectEager(query, key, value); + 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, 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); From 5bcc489c4f6f2573b7d7670803d31728463560dd Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Fri, 14 Aug 2026 21:51:19 +0800 Subject: [PATCH 03/21] fix: synchronize parallel linear fallback threads --- mllm/backends/cpu/ops/ParallelLinearOp.cpp | 3 +++ mllm/models/lfm2/modeling_lfm2.hpp | 4 ++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/mllm/backends/cpu/ops/ParallelLinearOp.cpp b/mllm/backends/cpu/ops/ParallelLinearOp.cpp index d8dbfed28..5fe0d2a78 100644 --- a/mllm/backends/cpu/ops/ParallelLinearOp.cpp +++ b/mllm/backends/cpu/ops/ParallelLinearOp.cpp @@ -110,6 +110,9 @@ void CPUParallelLinearOp::forward(const std::vector& inputs, std::vector const auto& input = inputs[0]; if (tryForwardSharedInputKaiM1(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); } diff --git a/mllm/models/lfm2/modeling_lfm2.hpp b/mllm/models/lfm2/modeling_lfm2.hpp index 5943026bd..1e1afc1ea 100644 --- a/mllm/models/lfm2/modeling_lfm2.hpp +++ b/mllm/models/lfm2/modeling_lfm2.hpp @@ -24,8 +24,8 @@ inline auto makeLfm2LinearOptions(int32_t in_channels, int32_t out_channels, boo .out_channels = out_channels, .bias = bias, .impl_type = impl_type, - // OnePlus 13T source-bound screening keeps I8MM prefill above 80 tok/s - // with six workers, while four workers avoid decode GEMV oversubscription. + // Source-bound OnePlus 13T screening selected six workers for I8MM + // prefill and four workers to avoid decode GEMV oversubscription. .kai_w4a32_decode_thread_cap = 4, .kai_w4a32_prefill_thread_cap = 6}; } From 0912ace7ab04f3067c9bdb81661de71100738918 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Sat, 15 Aug 2026 11:31:30 +0800 Subject: [PATCH 04/21] fix(cpu): address LFM2.5 runtime review findings --- .../cpu/ops/GroupedQueryAttentionOp.cpp | 59 +++++++++--- mllm/backends/cpu/ops/LinearOp.cpp | 15 ++- mllm/backends/cpu/ops/LinearOp.hpp | 2 + tests/cpu/Lfm2ConfigTest.cpp | 5 + tests/cpu/Qwen35GDNConvTest.cpp | 45 +++++---- tests/nn/GroupedQueryAttentionTest.cpp | 92 ++++++++++++++++++- 6 files changed, 179 insertions(+), 39 deletions(-) diff --git a/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp index 3705b3218..084aef9ad 100644 --- a/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp +++ b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp @@ -19,15 +19,25 @@ void groupedQueryAttentionDirectStrided(const Tensor& query, const Tensor& key, 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]; - - 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]); + // 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) { @@ -44,21 +54,43 @@ void groupedQueryAttentionDirectStrided(const Tensor& query, const Tensor& key, } } } + 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 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 += static_cast(query_rows[query_row][dim]) * static_cast(key_rows[key_row][dim]); + const Scalar* key_row = key_row_data[key_row_base + key_index]; + // Keep the legacy contiguous dot expression: changing it to a + // runtime-stride induction loop changes contraction/codegen and can + // diverge LFM2.5's frozen 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]); @@ -69,14 +101,15 @@ void groupedQueryAttentionDirectStrided(const Tensor& query, const Tensor& key, 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) { - const size_t value_row = (static_cast(batch) * v_shape[1] + kv_head) * v_shape[2] + key_index; - accumulated += (scores[key_index] / denominator) * static_cast(value_rows[value_row][value_dim]); + accumulated += + (scores[key_index] * static_cast(value_row_data[value_row_base + key_index][value_dim * v_stride[3]])) + * inverse_denominator; } - const size_t output_row = (static_cast(batch) * q_shape[1] + query_head) * q_shape[2] + query_index; - output_rows[output_row][value_dim] = static_cast(accumulated); + output_row[value_dim * o_stride[3]] = static_cast(accumulated); } } } diff --git a/mllm/backends/cpu/ops/LinearOp.cpp b/mllm/backends/cpu/ops/LinearOp.cpp index 4105c6273..61b6f31d6 100644 --- a/mllm/backends/cpu/ops/LinearOp.cpp +++ b/mllm/backends/cpu/ops/LinearOp.cpp @@ -253,10 +253,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: { @@ -293,10 +294,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: { @@ -311,10 +313,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: { @@ -331,10 +334,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: { @@ -349,10 +353,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 315a5dbfa..c45e582a2 100644 --- a/mllm/backends/cpu/ops/LinearOp.hpp +++ b/mllm/backends/cpu/ops/LinearOp.hpp @@ -14,6 +14,8 @@ constexpr bool shouldUseKaiW4A32I8mmPrefill(int m, bool disabled, bool cpu_suppo } constexpr int kaiW4A32ThreadCount(int m, int requested_threads, int decode_thread_cap, int prefill_thread_cap) { + // Every dynamic-input W4A32 KAI tile interprets the optional caps through + // this helper; zero keeps the repository-wide requested thread count. const int cap = m == 1 ? decode_thread_cap : prefill_thread_cap; return cap > 0 && cap < requested_threads ? cap : requested_threads; } diff --git a/tests/cpu/Lfm2ConfigTest.cpp b/tests/cpu/Lfm2ConfigTest.cpp index f7faff726..6cd23231a 100644 --- a/tests/cpu/Lfm2ConfigTest.cpp +++ b/tests/cpu/Lfm2ConfigTest.cpp @@ -56,4 +56,9 @@ TEST(Lfm2ConfigTest, KaiW4A32ThreadCapsSeparateDecodeAndPrefill) { EXPECT_EQ(kaiW4A32ThreadCount(1, 2, 4, 6), 2); EXPECT_EQ(kaiW4A32ThreadCount(28, 4, 4, 6), 4); EXPECT_EQ(kaiW4A32ThreadCount(1, 8, 0, 0), 8); + EXPECT_EQ(kaiW4A32ThreadCount(28, 8, 0, 0), 8); + EXPECT_EQ(kaiW4A32ThreadCount(1, 8, 4, 0), 4); + EXPECT_EQ(kaiW4A32ThreadCount(28, 8, 0, 6), 6); + EXPECT_EQ(kaiW4A32ThreadCount(1, 8, 12, 12), 8); + EXPECT_EQ(kaiW4A32ThreadCount(28, 8, 12, 12), 8); } diff --git a/tests/cpu/Qwen35GDNConvTest.cpp b/tests/cpu/Qwen35GDNConvTest.cpp index 9816e2053..9aafd5110 100644 --- a/tests/cpu/Qwen35GDNConvTest.cpp +++ b/tests/cpu/Qwen35GDNConvTest.cpp @@ -169,26 +169,31 @@ TEST(Qwen35GDNConvTest, HistoryFirstK3MatchesScalarReferenceBitwiseForLfmWidths) for (int batch : {1, 2}) { for (int sequence : {1, 2, 28, 225}) { for (int channels : {1, 3, 4, 5, 2045, 2048}) { - 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 = makeBuffer(state_count, 19); - - auto kernel_state = initial_state; - std::vector kernel_output(element_count, 0.0F); - 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); - - ASSERT_EQ(kernel_output, reference_output) - << "history-first output mismatch for B=" << batch << " S=" << sequence << " C=" << channels; - ASSERT_EQ(kernel_state, reference_state) - << "history-first state mismatch for B=" << batch << " S=" << sequence << " C=" << channels; + for (bool non_zero_history : {false, true}) { + 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); + 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); + + ASSERT_EQ(kernel_output, reference_output) + << "history-first output mismatch for B=" << batch << " S=" << sequence << " C=" << channels + << " history=" << (non_zero_history ? "non-zero" : "zero"); + ASSERT_EQ(kernel_state, reference_state) + << "history-first state mismatch for B=" << batch << " S=" << sequence << " C=" << channels + << " history=" << (non_zero_history ? "non-zero" : "zero"); + } } } } diff --git a/tests/nn/GroupedQueryAttentionTest.cpp b/tests/nn/GroupedQueryAttentionTest.cpp index 43ce947c0..ea29acad3 100644 --- a/tests/nn/GroupedQueryAttentionTest.cpp +++ b/tests/nn/GroupedQueryAttentionTest.cpp @@ -124,6 +124,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(); @@ -153,6 +220,29 @@ TEST_F(GroupedQueryAttentionTest, RegisteredDirectStridedMatchesReference) { 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), @@ -268,7 +358,7 @@ TEST_F(GroupedQueryAttentionTest, SupportsTransposedNonContiguousHeadViews) { auto key = key_bshd.transpose(1, 2); auto value = value_bshd.transpose(1, 2); - const auto actual = mllm::nn::llm_components::groupedQueryAttention(query, key, value); + const auto actual = mllm::nn::functional::groupedQueryAttention(query, key, value); const auto expected = gqaReference(query.contiguous(), key.contiguous(), value.contiguous()); expectNear(actual, expected); } From 5610f07db53d3cd2e2c0217d9748f3ae443486d1 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Fri, 21 Aug 2026 14:32:29 +0800 Subject: [PATCH 05/21] perf(cpu): extend KAI shared-input fusion to I8MM prefill Generalize matmul_shared_input_m1 into matmul_shared_input with an M dimension so gate/up and q/k/v projections share one packed LHS during prefill as well as decode. Work is distributed over the combined M x N tile grid with overflow guards, and matmul_shared_input_m1 stays as a thin wrapper. The product path remains gated on the existing I8MM prefill screen, so hosts without I8MM keep the per-projection fallback. Add two ARM benchmarks that compare the fused and independent paths and assert bitwise-equal sentinel hashes. --- benchmarks/cpu/CMakeLists.txt | 8 +- benchmarks/cpu/lfm2_parallel_linear.cpp | 295 ++++++++++++++++++ .../cpu/lfm2_parallel_linear_shared_mx.cpp | 265 ++++++++++++++++ mllm/backends/cpu/kernels/arm/linear/kai.cpp | 50 ++- mllm/backends/cpu/kernels/arm/linear/kai.hpp | 4 + mllm/backends/cpu/ops/LinearOp.cpp | 13 +- mllm/backends/cpu/ops/LinearOp.hpp | 2 + mllm/backends/cpu/ops/ParallelLinearOp.cpp | 47 ++- mllm/backends/cpu/ops/ParallelLinearOp.hpp | 4 +- 9 files changed, 652 insertions(+), 36 deletions(-) create mode 100644 benchmarks/cpu/lfm2_parallel_linear.cpp create mode 100644 benchmarks/cpu/lfm2_parallel_linear_shared_mx.cpp 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..b6c730269 --- /dev/null +++ b/benchmarks/cpu/lfm2_parallel_linear.cpp @@ -0,0 +1,295 @@ +// 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" + +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 long parsed = std::strtol(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 { + 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..eba69bdd7 --- /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 long parsed = std::strtol(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/mllm/backends/cpu/kernels/arm/linear/kai.cpp b/mllm/backends/cpu/kernels/arm/linear/kai.cpp index a5bf15689..37aab1e1d 100644 --- a/mllm/backends/cpu/kernels/arm/linear/kai.cpp +++ b/mllm/backends/cpu/kernels/arm/linear/kai.cpp @@ -445,10 +445,12 @@ void KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::matmul(float* __restrict__ dst, con } } -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) { - if (lhs_fp32 == nullptr || projections == nullptr || projection_count < 2 || workspace == nullptr || K <= 0 +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; } @@ -460,43 +462,61 @@ bool KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::matmul_shared_input_m1( } const auto& ukernel = ukernels_.at(tile_cfg); - kai_run_lhs_quant_pack_qai8dxp_f32(1, K, ukernel.get_mr(), ukernel.get_kr(), ukernel.get_sr(), 0, lhs_fp32, - K * sizeof(float), workspace); + 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) { - total_tiles += (static_cast(projections[projection_index].n) + n_step - 1) / n_step; + 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; } - const void* lhs_ptr = - static_cast(static_cast(workspace) + ukernel.get_lhs_packed_offset(0, K)); 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_tiles = (static_cast(projections[projection_index].n) + n_step - 1) / n_step; + 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 int n_index = static_cast(local_tile * n_step); + 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* rhs_ptr = static_cast( - reinterpret_cast(projection.packed_weight_bias) + ukernel.get_rhs_packed_offset(n_index, K, 32)); + 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(0, n_index, dst_stride)); + + ukernel.get_dst_offset(m_index, n_index, dst_stride)); - ukernel.run_matmul(1, actual_n, K, 32, lhs_ptr, rhs_ptr, dst_ptr, dst_stride, sizeof(float), + 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 e282bcc5c..bce54b9a4 100644 --- a/mllm/backends/cpu/kernels/arm/linear/kai.hpp +++ b/mllm/backends/cpu/kernels/arm/linear/kai.hpp @@ -124,6 +124,10 @@ 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); diff --git a/mllm/backends/cpu/ops/LinearOp.cpp b/mllm/backends/cpu/ops/LinearOp.cpp index 61b6f31d6..df0db22f8 100644 --- a/mllm/backends/cpu/ops/LinearOp.cpp +++ b/mllm/backends/cpu/ops/LinearOp.cpp @@ -43,8 +43,7 @@ bool cpuSupportsI8mm() { } KaiW4A32Tile selectKaiW4A32PrefillTile(int m) { - static const bool disabled = environmentFlagEnabled("MLLM_KAI_PREFILL_I8MM_DISABLE"); - if (detail::shouldUseKaiW4A32I8mmPrefill(m, disabled, cpuSupportsI8mm())) { return kKaiW4A32I8mmTile; } + if (detail::shouldUseKaiW4A32I8mmPrefill(m)) { return kKaiW4A32I8mmTile; } return kKaiW4A32DotProdTile; } @@ -70,6 +69,16 @@ void traceKaiW4A32PrefillTile(KaiW4A32Tile tile, int m, int k, int n, int thread } // namespace +bool detail::shouldUseKaiW4A32I8mmPrefill(int m) { +#if defined(MLLM_HOST_ARCH_ARM64) || defined(MLLM_HOST_ARCH_ARM) + static const bool disabled = environmentFlagEnabled("MLLM_KAI_PREFILL_I8MM_DISABLE"); + return shouldUseKaiW4A32I8mmPrefill(m, disabled, cpuSupportsI8mm()); +#else + (void)m; + return false; +#endif +} + CPULinearOp::CPULinearOp(const aops::LinearOpOptions& options) : LinearOp(options) {} int CPULinearOp::kaiW4A32ThreadCount(int m) const { diff --git a/mllm/backends/cpu/ops/LinearOp.hpp b/mllm/backends/cpu/ops/LinearOp.hpp index c45e582a2..46e4df87e 100644 --- a/mllm/backends/cpu/ops/LinearOp.hpp +++ b/mllm/backends/cpu/ops/LinearOp.hpp @@ -13,6 +13,8 @@ constexpr bool shouldUseKaiW4A32I8mmPrefill(int m, bool disabled, bool cpu_suppo return m >= 4 && !disabled && cpu_supports_i8mm; } +bool shouldUseKaiW4A32I8mmPrefill(int m); + constexpr int kaiW4A32ThreadCount(int m, int requested_threads, int decode_thread_cap, int prefill_thread_cap) { // Every dynamic-input W4A32 KAI tile interprets the optional caps through // this helper; zero keeps the repository-wide requested thread count. diff --git a/mllm/backends/cpu/ops/ParallelLinearOp.cpp b/mllm/backends/cpu/ops/ParallelLinearOp.cpp index 5fe0d2a78..3b8cc0c43 100644 --- a/mllm/backends/cpu/ops/ParallelLinearOp.cpp +++ b/mllm/backends/cpu/ops/ParallelLinearOp.cpp @@ -37,35 +37,49 @@ void CPUParallelLinearOp::load(const ParameterFile::ptr_t& ploader) { } Tensor CPUParallelLinearOp::acquireKaiWorkspace(int32_t workspace_size) { - if (kai_decode_workspace_.isNil() || kai_decode_workspace_.numel() < static_cast(workspace_size)) { - kai_decode_workspace_ = Tensor::empty({workspace_size}, kInt8, kCPU).alloc(); + if (kai_workspace_.isNil() || kai_workspace_.numel() < static_cast(workspace_size)) { + kai_workspace_ = Tensor::empty({workspace_size}, kInt8, kCPU).alloc(); } - return kai_decode_workspace_; + return kai_workspace_; } -bool CPUParallelLinearOp::tryForwardSharedInputKaiM1(const Tensor& input, std::vector& outputs) { +bool CPUParallelLinearOp::tryForwardSharedInputKai(const Tensor& input, std::vector& outputs) { #if defined(MLLM_HOST_ARCH_ARM64) || defined(MLLM_HOST_ARCH_ARM) constexpr size_t kMaximumSharedProjections = 3; constexpr auto kRequiredImpl = aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32; using KaiHelper = ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk; - constexpr auto kTile = KaiHelper::Tiles::qai8dxp1x8_qsi4c32p8x8_1x8x32; + constexpr auto kDecodeTile = KaiHelper::Tiles::qai8dxp1x8_qsi4c32p8x8_1x8x32; + constexpr auto kPrefillTile = KaiHelper::Tiles::qai8dxp4x8_qsi4c32p8x8_4x8x32; if (input.isNil() || input.device() != kCPU || input.dtype() != kFloat32 || !input.isContiguous() || input.rank() < 2 - || input.size(-2) != 1 || input.size(-1) != options_.in_channels || options_.bias || options_.impl_type != kRequiredImpl + || 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; } + const auto tile = m == 1 ? kDecodeTile : kPrefillTile; + if (m > 1 && !detail::shouldUseKaiW4A32I8mmPrefill(m)) { + // The generic shared-input dot-product prefill path has not passed the + // mobile product screen; preserve the established per-projection fallback. + return false; + } for (size_t index = 0; index < weights_.size(); ++index) { - if (weights_[index].isNil() || weights_[index].device() != kCPU || outputs[index].dtype() != kFloat32 - || outputs[index].device() != kCPU) { + 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 int32_t thread_count = detail::kaiW4A32ThreadCount(1, options_.getThreads(), options_.kai_w4a32_decode_thread_cap, + const int32_t thread_count = detail::kaiW4A32ThreadCount(m, options_.getThreads(), options_.kai_w4a32_decode_thread_cap, options_.kai_w4a32_prefill_thread_cap); std::array projections{}; for (size_t index = 0; index < weights_.size(); ++index) { @@ -77,11 +91,11 @@ bool CPUParallelLinearOp::tryForwardSharedInputKaiM1(const Tensor& input, std::v } KaiHelper kai_helper; - const size_t workspace_size = kai_helper.workspace_size(1, options_.in_channels, kTile); + const size_t workspace_size = kai_helper.workspace_size(m, options_.in_channels, tile); if (workspace_size == 0 || workspace_size > static_cast(std::numeric_limits::max())) { return false; } auto workspace = acquireKaiWorkspace(static_cast(workspace_size)); - if (!kai_helper.matmul_shared_input_m1(input.ptr(), projections.data(), weights_.size(), workspace.ptr(), - options_.in_channels, kTile, thread_count)) { + if (!kai_helper.matmul_shared_input(input.ptr(), projections.data(), weights_.size(), workspace.ptr(), m, + options_.in_channels, tile, thread_count)) { return false; } @@ -90,12 +104,13 @@ bool CPUParallelLinearOp::tryForwardSharedInputKaiM1(const Tensor& input, std::v return value != nullptr && value[0] == '1' && value[1] == '\0'; }(); if (trace_activation) { - const uint32_t activation_bit = 1U << static_cast(weights_.size() - 2); + const uint32_t projection_group = static_cast(weights_.size() - 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 k=%d threads=%d\n", weights_.size(), options_.in_channels, - thread_count); + std::fprintf(stderr, "MLLM_KAI_SHARED_INPUT_ACTIVATED rhs=%zu m=%d k=%d threads=%d tile=%s\n", weights_.size(), m, + options_.in_channels, thread_count, m == 1 ? "dotprod_1x8" : "i8mm_4x8"); } } return true; @@ -108,7 +123,7 @@ bool CPUParallelLinearOp::tryForwardSharedInputKaiM1(const Tensor& input, std::v void CPUParallelLinearOp::forward(const std::vector& inputs, std::vector& outputs) { const auto& input = inputs[0]; - if (tryForwardSharedInputKaiM1(input, outputs)) { return; } + 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. diff --git a/mllm/backends/cpu/ops/ParallelLinearOp.hpp b/mllm/backends/cpu/ops/ParallelLinearOp.hpp index 14a6faba6..394611f1c 100644 --- a/mllm/backends/cpu/ops/ParallelLinearOp.hpp +++ b/mllm/backends/cpu/ops/ParallelLinearOp.hpp @@ -19,11 +19,11 @@ class CPUParallelLinearOp final : public aops::ParallelLinearOp { void forward(const std::vector& inputs, std::vector& outputs) override; private: - bool tryForwardSharedInputKaiM1(const Tensor& input, std::vector& outputs); + bool tryForwardSharedInputKai(const Tensor& input, std::vector& outputs); Tensor acquireKaiWorkspace(int32_t workspace_size); std::vector> fallback_ops_; - Tensor kai_decode_workspace_; + Tensor kai_workspace_; }; class CPUParallelLinearOpFactory : public TypedOpFactory { From 1e48f4542d48ccf9780d44fdda3d713a95bc1a74 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Fri, 21 Aug 2026 15:08:02 +0800 Subject: [PATCH 06/21] refactor(cpu): make the causal conv operation target-neutral Move the history-first depthwise causal convolution out of the gated delta net directory into kernels/common/causal_conv, where a reusable causal-convolution primitive belongs, and give it its own focused bitwise oracle instead of hosting it in the Qwen3.5 GDN test. Rename the activation hook from MLLM_LFM2_SHORT_CONV_TRACE to MLLM_CAUSAL_CONV1D_TRACE and emit one marker per accumulation order, so a framework-level operation no longer reports under a single model's name. Drop the remaining model-specific wording from the shared kernel and grouped-query attention comments. --- .../causal_conv/depthwise_causal_conv.cpp | 71 ++++++++++++++ .../causal_conv/depthwise_causal_conv.hpp | 18 ++++ .../kernels/common/gdn/gated_delta_net.cpp | 56 ----------- .../kernels/common/gdn/gated_delta_net.hpp | 6 -- .../cpu/ops/CausalDepthwiseConv1DOp.cpp | 38 +++++--- .../cpu/ops/GroupedQueryAttentionOp.cpp | 6 +- tests/cpu/CMakeLists.txt | 4 + tests/cpu/CausalDepthwiseConvKernelTest.cpp | 94 +++++++++++++++++++ tests/cpu/Qwen35GDNConvTest.cpp | 59 ------------ 9 files changed, 213 insertions(+), 139 deletions(-) create mode 100644 mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.cpp create mode 100644 mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp create mode 100644 tests/cpu/CausalDepthwiseConvKernelTest.cpp 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/gdn/gated_delta_net.cpp b/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.cpp index db279cf26..b20e6fc3b 100644 --- a/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.cpp +++ b/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.cpp @@ -243,62 +243,6 @@ void depthwiseCausalConvF32(const float* input, const float* weight, float* stat } } -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__) - // LFM2.5 uses K=3. Four adjacent channels are deinterleaved 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]; - } - } - } -} - void gatedDeltaRuleF32(const float* q, const float* k, const float* v, const float* a, const float* b, const float* a_log, const float* dt_bias, float* state, float* output, int batch_size, int sequence_length, int num_key_heads, int num_value_heads, int key_head_dim, int value_head_dim) { diff --git a/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp b/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp index e95171e42..78a53ac81 100644 --- a/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp +++ b/mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp @@ -10,12 +10,6 @@ namespace mllm::cpu::gdn { void depthwiseCausalConvF32(const float* input, const float* weight, float* state, float* output, int batch_size, int sequence_length, int channels, int kernel_size); -// Same [B, S, C] / [B, C, K - 1] state contract, with the accumulation order -// used by CPUConv1D: zero, historical taps in ascending order, then current. -// Keeping this explicit lets callers retain bitwise-sensitive model semantics. -void depthwiseCausalConvHistoryFirstF32(const float* input, const float* weight, float* state, float* output, - int batch_size, int sequence_length, int channels, int kernel_size); - // Gated-delta recurrence. Each [batch, value_head] state is independent and // may run in parallel, while tokens within one state remain strictly ordered. // diff --git a/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.cpp b/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.cpp index 1a7b10df0..03de5bb4a 100644 --- a/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.cpp +++ b/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.cpp @@ -4,11 +4,13 @@ #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 { @@ -31,26 +33,32 @@ void CPUCausalDepthwiseConv1DOp::forward(const std::vector& inputs, std: auto& updated_state = outputs[1]; if (!options_.state_inplace) { std::memcpy(updated_state.ptr(), state.ptr(), state.bytes()); } - if (options_.accumulation_order == aops::CausalDepthwiseConv1DAccumulationOrder::kHistoryFirst) { - gdn::depthwiseCausalConvHistoryFirstF32(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_LFM2_SHORT_CONV_TRACE"); - return value != nullptr && value[0] == '1' && value[1] == '\0'; - }(); - if (trace_activation) { - static std::atomic activated{false}; - if (!activated.exchange(true, std::memory_order_relaxed)) { - std::fprintf(stderr, "MLLM_LFM2_SHORT_CONV_REUSE_ACTIVATED k=%d channels=%d\n", options_.kernel_size, - options_.channels); - } - } + 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}) { diff --git a/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp index 084aef9ad..c20b58b2c 100644 --- a/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp +++ b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp @@ -76,9 +76,9 @@ void groupedQueryAttentionDirectStrided(const Tensor& query, const Tensor& key, 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 legacy contiguous dot expression: changing it to a - // runtime-stride induction loop changes contraction/codegen and can - // diverge LFM2.5's frozen generation-token oracle. + // 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]); diff --git a/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index aae1210df..b028c234d 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -10,6 +10,10 @@ add_executable(Mllm-Test-Qwen35-GDN-Conv Qwen35GDNConvTest.cpp) target_link_libraries(Mllm-Test-Qwen35-GDN-Conv PRIVATE gtest_main MllmCPUBackend) target_include_directories(Mllm-Test-Qwen35-GDN-Conv PRIVATE ${MLLM_INCLUDE_DIR}) +add_executable(Mllm-Test-CausalDepthwiseConvKernel CausalDepthwiseConvKernelTest.cpp) +target_link_libraries(Mllm-Test-CausalDepthwiseConvKernel PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-CausalDepthwiseConvKernel PRIVATE ${MLLM_INCLUDE_DIR}) + add_executable(Mllm-Test-KaiW4A32Pack KaiW4A32PackTest.cpp) target_link_libraries(Mllm-Test-KaiW4A32Pack PRIVATE gtest_main MllmCPUBackend) target_include_directories(Mllm-Test-KaiW4A32Pack PRIVATE ${MLLM_INCLUDE_DIR}) diff --git a/tests/cpu/CausalDepthwiseConvKernelTest.cpp b/tests/cpu/CausalDepthwiseConvKernelTest.cpp new file mode 100644 index 000000000..e40823707 --- /dev/null +++ b/tests/cpu/CausalDepthwiseConvKernelTest.cpp @@ -0,0 +1,94 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +// Focused oracle for the history-first depthwise causal convolution kernel. +// +// The reference below is an independent scalar implementation of the frozen +// contract. It is deliberately not routed through the production kernel, so a +// vectorized fast path cannot validate itself. Both the output and the final +// history are compared bitwise: an output-only comparison would miss a +// corrupted history that only shows up in the next chunk. + +#include + +#include +#include + +#include "mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp" + +namespace { + +using mllm::cpu::causal_conv::depthwiseCausalConvHistoryFirstF32; + +// Deterministic index-derived fill. No RNG, so every host reproduces the same +// bytes without carrying a seed through the evidence record. +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; +} + +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; +} + +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]; + } + } + } +} + +TEST(CausalDepthwiseConvKernelTest, HistoryFirstK3MatchesScalarReferenceBitwise) { + constexpr int kKernel = 3; + for (int batch : {1, 2}) { + for (int sequence : {1, 2, 28, 225}) { + for (int channels : {1, 3, 4, 5, 2045, 2048}) { + for (bool non_zero_history : {false, true}) { + 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); + 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); + + ASSERT_EQ(kernel_output, reference_output) + << "history-first output mismatch for B=" << batch << " S=" << sequence << " C=" << channels + << " history=" << (non_zero_history ? "non-zero" : "zero"); + ASSERT_EQ(kernel_state, reference_state) + << "history-first state mismatch for B=" << batch << " S=" << sequence << " C=" << channels + << " history=" << (non_zero_history ? "non-zero" : "zero"); + } + } + } + } +} + +} // namespace diff --git a/tests/cpu/Qwen35GDNConvTest.cpp b/tests/cpu/Qwen35GDNConvTest.cpp index 9aafd5110..46bcaa079 100644 --- a/tests/cpu/Qwen35GDNConvTest.cpp +++ b/tests/cpu/Qwen35GDNConvTest.cpp @@ -23,7 +23,6 @@ namespace { using mllm::cpu::gdn::depthwiseCausalConvF32; -using mllm::cpu::gdn::depthwiseCausalConvHistoryFirstF32; // Deterministic index-derived fill. No RNG, so every host reproduces the same // bytes without carrying a seed through the evidence record. @@ -62,29 +61,6 @@ void referenceDepthwiseCausalConv(const std::vector& input, const std::ve } } -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]; - } - } - } -} - struct ConvCase { int batch; int sequence; @@ -164,41 +140,6 @@ TEST(Qwen35GDNConvTest, MatchesScalarReferenceWithChannelTailAtProductionScale) } } -TEST(Qwen35GDNConvTest, HistoryFirstK3MatchesScalarReferenceBitwiseForLfmWidths) { - constexpr int kKernel = 3; - for (int batch : {1, 2}) { - for (int sequence : {1, 2, 28, 225}) { - for (int channels : {1, 3, 4, 5, 2045, 2048}) { - for (bool non_zero_history : {false, true}) { - 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); - 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); - - ASSERT_EQ(kernel_output, reference_output) - << "history-first output mismatch for B=" << batch << " S=" << sequence << " C=" << channels - << " history=" << (non_zero_history ? "non-zero" : "zero"); - ASSERT_EQ(kernel_state, reference_state) - << "history-first state mismatch for B=" << batch << " S=" << sequence << " C=" << channels - << " history=" << (non_zero_history ? "non-zero" : "zero"); - } - } - } - } -} - TEST(Qwen35GDNConvTest, ChunkedPartitionsMatchOneShot) { struct Partition { int channels; From 7fa7ea65a84bb5b36eec1e333796c4f984339b25 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Fri, 21 Aug 2026 15:14:44 +0800 Subject: [PATCH 07/21] refactor(models): share default RoPE tables and guard fused projection names Lift the plain default-RoPE inverse-frequency and sin/cos table helpers out of the LFM2 and MiniCPM5 model headers into one shared model-side header, with the input validation the model-local copies never had. Both models used the identical no-scaling variant, so the generated tables are unchanged. The helpers stay under mllm/models because they materialize constant operation inputs rather than performing tensor computation, which nn/llm_components must not host. ParallelLinear resolves parameters in its parent scope to keep original checkpoint names, which makes ambiguous projection names bind the wrong tensors. Reject duplicate, empty, and scope-escaping names in reshape and load, and document why the operation's own name is not part of the parameter path. --- mllm/core/aops/ParallelLinearOp.cpp | 28 ++++++-- mllm/core/aops/ParallelLinearOp.hpp | 2 + mllm/models/common/rope_tables.hpp | 79 ++++++++++++++++++++++ mllm/models/lfm2/modeling_lfm2.hpp | 41 ++--------- mllm/models/minicpm5/modeling_minicpm5.hpp | 42 ++---------- tests/cpu/Lfm2RegisteredOpsTest.cpp | 25 +++++++ 6 files changed, 142 insertions(+), 75 deletions(-) create mode 100644 mllm/models/common/rope_tables.hpp diff --git a/mllm/core/aops/ParallelLinearOp.cpp b/mllm/core/aops/ParallelLinearOp.cpp index 26142cd9c..48e3f3824 100644 --- a/mllm/core/aops/ParallelLinearOp.cpp +++ b/mllm/core/aops/ParallelLinearOp.cpp @@ -16,16 +16,36 @@ 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::load(const ParameterFile::ptr_t& ploader) { +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()); @@ -76,10 +96,8 @@ void ParallelLinearOp::reshape(const std::vector& inputs, std::vector weights_; std::vector biases_; ParallelLinearOpOptions options_; 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/modeling_lfm2.hpp b/mllm/models/lfm2/modeling_lfm2.hpp index 1e1afc1ea..170472375 100644 --- a/mllm/models/lfm2/modeling_lfm2.hpp +++ b/mllm/models/lfm2/modeling_lfm2.hpp @@ -10,6 +10,7 @@ #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" @@ -18,6 +19,11 @@ 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 auto makeLfm2LinearOptions(int32_t in_channels, int32_t out_channels, bool bias, aops::LinearImplTypes impl_type) -> aops::LinearOpOptions { return {.in_channels = in_channels, @@ -42,41 +48,6 @@ inline auto makeLfm2ParallelLinearOptions(int32_t in_channels, std::vector Tensor { - 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; -} - -inline auto makeRotaryPosEmbedding(const Tensor& position_ids, const Tensor& inv_freq) -> std::pair { - 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}; -} - class Lfm2MLP final : public nn::Module { public: Lfm2MLP() = default; diff --git a/mllm/models/minicpm5/modeling_minicpm5.hpp b/mllm/models/minicpm5/modeling_minicpm5.hpp index a405ebb32..a15aa08ba 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: @@ -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/tests/cpu/Lfm2RegisteredOpsTest.cpp b/tests/cpu/Lfm2RegisteredOpsTest.cpp index 480d6e09c..0d0cee21a 100644 --- a/tests/cpu/Lfm2RegisteredOpsTest.cpp +++ b/tests/cpu/Lfm2RegisteredOpsTest.cpp @@ -5,10 +5,14 @@ #include #include +#include +#include #include +#include #include #include "mllm/compile/ir/Trace.hpp" +#include "mllm/core/aops/ParallelLinearOp.hpp" #include "mllm/compile/ir/linalg/Op.hpp" #include "mllm/compile/jit/binary/LinalgIRSerialization.hpp" #include "mllm/compile/jit/interpreter/AopsFromJson.hpp" @@ -132,4 +136,25 @@ TEST_F(Lfm2RegisteredOpsTest, ParallelLinearTracesAndSerializesProjectionContrac EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kParallelLinear); } +// Fused parameters resolve in the parent scope, so ambiguous or scope-escaping +// projection names would silently bind the wrong checkpoint tensors. +TEST_F(Lfm2RegisteredOpsTest, ParallelLinearRejectsAmbiguousProjectionNames) { + auto reshapeWith = [](std::vector projection_names) { + auto op = std::make_shared( + mllm::aops::ParallelLinearOpOptions{.in_channels = 2, + .out_channels = {2, 1}, + .projection_names = std::move(projection_names), + .bias = false, + .impl_type = mllm::aops::LinearImplTypes::kGGUF}); + std::vector inputs = {Tensor::empty({1, 1, 2}, mllm::kFloat32, mllm::kCPU)}; + std::vector outputs; + op->reshape(inputs, outputs); + }; + + EXPECT_THROW(reshapeWith({"same", "same"}), std::invalid_argument); + EXPECT_THROW(reshapeWith({"left", "nested.right"}), std::invalid_argument); + EXPECT_THROW(reshapeWith({"left", ""}), std::invalid_argument); + EXPECT_NO_THROW(reshapeWith({"left", "right"})); +} + } // namespace From 973c1499eb6f88aa95cfe289ce57275fa3909cfe Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Fri, 21 Aug 2026 15:38:57 +0800 Subject: [PATCH 08/21] refactor: converge grouped-query attention into one operation Decode-only grouped-query attention and the general path were two framework operations with overlapping semantics, so a new model had no way to tell which one it should reach for. Fold the decode operation into GroupedQueryAttention as the DecodeNativeKV implementation: it keeps its own reduction order and single-query-position contract, and still runs the same decode kernel, so MiniCPM5 generation is unchanged. nn::functional::groupedQueryAttentionDecode stays as the public entry point. Graphs serialized under the old "GroupedQueryAttentionDecode" op type still reconstruct, and OpTypes value 76 is retired rather than reused so an old graph can never alias a different operation. --- mllm/backends/cpu/CPUBackend.cpp | 3 +- .../cpu/ops/GroupedQueryAttentionDecodeOp.cpp | 104 ------------------ .../cpu/ops/GroupedQueryAttentionDecodeOp.hpp | 26 ----- .../cpu/ops/GroupedQueryAttentionOp.cpp | 102 +++++++++++++++-- mllm/compile/ir/GeneratedRTTIKind.hpp | 3 +- mllm/compile/ir/NodeRTTIClassOfImpl.hpp | 6 +- mllm/compile/ir/linalg/Op.cpp | 1 - mllm/compile/ir/linalg/Op.hpp | 2 - mllm/compile/ir/rtti_kind_gen.py | 1 - .../jit/binary/LinalgIRSerialization.cpp | 3 - .../jit/binary/LinalgIRSerialization.hpp | 1 - mllm/compile/jit/interpreter/AopsFromJson.cpp | 9 +- mllm/core/OpTypes.hpp | 6 +- .../aops/GroupedQueryAttentionDecodeOp.cpp | 63 ----------- .../aops/GroupedQueryAttentionDecodeOp.hpp | 36 ------ mllm/core/aops/GroupedQueryAttentionOp.cpp | 8 ++ mllm/core/aops/GroupedQueryAttentionOp.hpp | 11 ++ mllm/models/minicpm5/modeling_minicpm5.hpp | 6 +- mllm/nn/Functional.cpp | 4 +- mllm/nn/Nn.hpp | 1 - .../nn/layers/GroupedQueryAttentionDecode.cpp | 14 --- .../nn/layers/GroupedQueryAttentionDecode.hpp | 20 ---- tests/nn/GroupedQueryAttentionTest.cpp | 32 +++--- 23 files changed, 146 insertions(+), 316 deletions(-) delete mode 100644 mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.cpp delete mode 100644 mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.hpp delete mode 100644 mllm/core/aops/GroupedQueryAttentionDecodeOp.cpp delete mode 100644 mllm/core/aops/GroupedQueryAttentionDecodeOp.hpp delete mode 100644 mllm/nn/layers/GroupedQueryAttentionDecode.cpp delete mode 100644 mllm/nn/layers/GroupedQueryAttentionDecode.hpp diff --git a/mllm/backends/cpu/CPUBackend.cpp b/mllm/backends/cpu/CPUBackend.cpp index db22e417c..5b59e7c91 100644 --- a/mllm/backends/cpu/CPUBackend.cpp +++ b/mllm/backends/cpu/CPUBackend.cpp @@ -25,7 +25,6 @@ #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" @@ -87,7 +86,7 @@ CPUBackend::CPUBackend() : Backend(kCPU, createCPUAllocator()) { CPUConv2DOpFactory, CPULayerNorm2DOpFactory, CPUInterpolateOpFactory, CPUPadOpFactory, CPUMaskedScatterOpFactory, CPUArgsortOpFactory, CPUCloneOpFactory, CPUAvgPool1dOpFactory, CPUFlashAttention2SwaSinkOpFactory, CPURadixAttnRelaxOpFactory, CPURadixAttnSwaSinkOpFactory, CPUEqualOpFactory, CPUWhereOpFactory, - CPUGatherOpFactory, CPUGroupedQueryAttentionDecodeOpFactory, CPUCausalDepthwiseConv1DOpFactory, + CPUGatherOpFactory, CPUCausalDepthwiseConv1DOpFactory, CPUGroupedQueryAttentionOpFactory, CPUParallelLinearOpFactory>(); } 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 index c20b58b2c..ba062692e 100644 --- a/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp +++ b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp @@ -9,6 +9,7 @@ #include #include +#include "mllm/backends/cpu/kernels/common/gqa_decode/fwd_bhsd.hpp" #include "mllm/core/Parallel.hpp" namespace mllm::cpu { @@ -116,20 +117,107 @@ void groupedQueryAttentionDirectStrided(const Tensor& query, const Tensor& key, 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(); + 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; + } + } + } +} + +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) { - if (options_.implementation != aops::GroupedQueryAttentionImplementation::kDirectStrided) { - throw std::invalid_argument("Unsupported CPU GroupedQueryAttention implementation"); - } - if (inputs[0].dtype() == kFloat32) { - groupedQueryAttentionDirectStrided(inputs[0], inputs[1], inputs[2], outputs[0]); - } else { - groupedQueryAttentionDirectStrided(inputs[0], inputs[1], inputs[2], outputs[0]); + 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/compile/ir/GeneratedRTTIKind.hpp b/mllm/compile/ir/GeneratedRTTIKind.hpp index 50e9b21cd..726ec8657 100644 --- a/mllm/compile/ir/GeneratedRTTIKind.hpp +++ b/mllm/compile/ir/GeneratedRTTIKind.hpp @@ -1,4 +1,4 @@ -// Auto generated: 2026-08-14 18:50:28 +// Auto generated: 2026-08-21 15:20:12 // do not modify this file #pragma once @@ -41,7 +41,6 @@ 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, diff --git a/mllm/compile/ir/NodeRTTIClassOfImpl.hpp b/mllm/compile/ir/NodeRTTIClassOfImpl.hpp index fd59ffe7f..0a08e255b 100644 --- a/mllm/compile/ir/NodeRTTIClassOfImpl.hpp +++ b/mllm/compile/ir/NodeRTTIClassOfImpl.hpp @@ -1,4 +1,4 @@ -// Auto generated: 2026-08-14 18:50:28 +// Auto generated: 2026-08-21 15:20:12 // do not modify this file #pragma once namespace mllm::ir { @@ -93,10 +93,6 @@ 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 diff --git a/mllm/compile/ir/linalg/Op.cpp b/mllm/compile/ir/linalg/Op.cpp index 150d71c0b..a33cd581b 100644 --- a/mllm/compile/ir/linalg/Op.cpp +++ b/mllm/compile/ir/linalg/Op.cpp @@ -66,7 +66,6 @@ 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); diff --git a/mllm/compile/ir/linalg/Op.hpp b/mllm/compile/ir/linalg/Op.hpp index 09d6831d3..e79765cc7 100644 --- a/mllm/compile/ir/linalg/Op.hpp +++ b/mllm/compile/ir/linalg/Op.hpp @@ -35,7 +35,6 @@ class X2XOp; class ViewOp; class SplitOp; class FlashAttention2Op; -class GroupedQueryAttentionDecodeOp; class CausalDepthwiseConv1DOp; class GroupedQueryAttentionOp; class ParallelLinearOp; @@ -201,7 +200,6 @@ 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); diff --git a/mllm/compile/ir/rtti_kind_gen.py b/mllm/compile/ir/rtti_kind_gen.py index e7b5c0ace..03029b23b 100644 --- a/mllm/compile/ir/rtti_kind_gen.py +++ b/mllm/compile/ir/rtti_kind_gen.py @@ -247,7 +247,6 @@ 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")) diff --git a/mllm/compile/jit/binary/LinalgIRSerialization.cpp b/mllm/compile/jit/binary/LinalgIRSerialization.cpp index f3197fe81..307556400 100644 --- a/mllm/compile/jit/binary/LinalgIRSerialization.cpp +++ b/mllm/compile/jit/binary/LinalgIRSerialization.cpp @@ -11,7 +11,6 @@ #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" @@ -73,7 +72,6 @@ nlohmann::json dumpLinalgIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { CASE(Split) CASE(STFT) CASE(FlashAttention2) - CASE(GroupedQueryAttentionDecode) CASE(CausalDepthwiseConv1D) CASE(GroupedQueryAttention) CASE(ParallelLinear) @@ -249,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 705b6f3e3..a3c2901d3 100644 --- a/mllm/compile/jit/binary/LinalgIRSerialization.hpp +++ b/mllm/compile/jit/binary/LinalgIRSerialization.hpp @@ -37,7 +37,6 @@ 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); diff --git a/mllm/compile/jit/interpreter/AopsFromJson.cpp b/mllm/compile/jit/interpreter/AopsFromJson.cpp index 449b2c542..f6c5c7a78 100644 --- a/mllm/compile/jit/interpreter/AopsFromJson.cpp +++ b/mllm/compile/jit/interpreter/AopsFromJson.cpp @@ -22,7 +22,6 @@ #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" @@ -613,13 +612,17 @@ 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) { diff --git a/mllm/core/OpTypes.hpp b/mllm/core/OpTypes.hpp index fb88f67c1..0a3263003 100644 --- a/mllm/core/OpTypes.hpp +++ b/mllm/core/OpTypes.hpp @@ -97,8 +97,9 @@ 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, @@ -189,7 +190,6 @@ 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"; 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 index fd0df3631..a41b19411 100644 --- a/mllm/core/aops/GroupedQueryAttentionOp.cpp +++ b/mllm/core/aops/GroupedQueryAttentionOp.cpp @@ -48,6 +48,14 @@ void GroupedQueryAttentionOp::reshape(const std::vector& inputs, std::ve 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())); } diff --git a/mllm/core/aops/GroupedQueryAttentionOp.hpp b/mllm/core/aops/GroupedQueryAttentionOp.hpp index e6aa99d95..884818775 100644 --- a/mllm/core/aops/GroupedQueryAttentionOp.hpp +++ b/mllm/core/aops/GroupedQueryAttentionOp.hpp @@ -11,19 +11,30 @@ 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); } diff --git a/mllm/models/minicpm5/modeling_minicpm5.hpp b/mllm/models/minicpm5/modeling_minicpm5.hpp index a15aa08ba..256f8d10a 100644 --- a/mllm/models/minicpm5/modeling_minicpm5.hpp +++ b/mllm/models/minicpm5/modeling_minicpm5.hpp @@ -63,7 +63,9 @@ 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::GroupedQueryAttentionOpOptions{.implementation = aops::GroupedQueryAttentionImplementation::kDecodeNativeKV}); } std::vector forward(const std::vector& inputs, const std::vector& args) override { @@ -95,7 +97,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; diff --git a/mllm/nn/Functional.cpp b/mllm/nn/Functional.cpp index db70e6d57..7ad399c67 100644 --- a/mllm/nn/Functional.cpp +++ b/mllm/nn/Functional.cpp @@ -6,7 +6,6 @@ #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" @@ -90,8 +89,7 @@ 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, diff --git a/mllm/nn/Nn.hpp b/mllm/nn/Nn.hpp index 49c653679..a1492d9ae 100644 --- a/mllm/nn/Nn.hpp +++ b/mllm/nn/Nn.hpp @@ -15,7 +15,6 @@ #include "mllm/nn/layers/GELU.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/GroupedQueryAttentionDecode.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 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/tests/nn/GroupedQueryAttentionTest.cpp b/tests/nn/GroupedQueryAttentionTest.cpp index ea29acad3..af66bb718 100644 --- a/tests/nn/GroupedQueryAttentionTest.cpp +++ b/tests/nn/GroupedQueryAttentionTest.cpp @@ -46,18 +46,6 @@ class GroupedQueryAttentionTraceModule 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_(); - } - 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; } - } - } - return nullptr; -} mllm::ir::linalg::GroupedQueryAttentionOp::ptr_t findGroupedQueryAttentionOp(const mllm::ir::node_ptr_t& node) { if (node->isa_()) { @@ -337,17 +325,27 @@ TEST_F(GroupedQueryAttentionTest, DecodeOpTraceAndSerializationRoundTrip) { 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) { From ae5a8672b410629bece378ef6b27f836267d30db Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Fri, 21 Aug 2026 16:48:41 +0800 Subject: [PATCH 09/21] fix(cpu): honour the output stride in the GQA decode fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The decode kernel declines any tensor whose last-dimension stride is not 1, so the scalar fallback is reached precisely when a non-unit output stride is possible — yet it indexed the output as if the value dimension were contiguous. Multiply by the output stride, matching what the DirectStrided path already does and what the kernel is handed. Not reachable today: reshape allocates the output through Tensor::empty, so the stride is 1 and the emitted addresses are unchanged. This keeps the fallback correct for any caller that supplies a strided output. --- mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp index ba062692e..a9cc82b0e 100644 --- a/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp +++ b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp @@ -127,6 +127,10 @@ void groupedQueryAttentionDecodeFloat32Reference(const Tensor& query, const Tens 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])); @@ -171,7 +175,7 @@ void groupedQueryAttentionDecodeFloat32Reference(const Tensor& query, const Tens 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; + output_head[static_cast(value_dim) * o_stride[3]] = result; } } } From 389710aca761735c5c078c4493f98aa6a42c7868 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Fri, 21 Aug 2026 16:48:41 +0800 Subject: [PATCH 10/21] fix(benchmark): parse sizes with a fixed-width integer type .clang-tidy enables google-* with WarningsAsErrors '*', so the plain long from std::strtol trips google-runtime-int and fails the build. Parse with int64_t and std::strtoll in both benchmark drivers; the range guard and return type are unchanged. --- benchmarks/cpu/lfm2_parallel_linear.cpp | 2 +- benchmarks/cpu/lfm2_parallel_linear_shared_mx.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmarks/cpu/lfm2_parallel_linear.cpp b/benchmarks/cpu/lfm2_parallel_linear.cpp index b6c730269..5320bd266 100644 --- a/benchmarks/cpu/lfm2_parallel_linear.cpp +++ b/benchmarks/cpu/lfm2_parallel_linear.cpp @@ -59,7 +59,7 @@ float deterministicValue(uint32_t& state) { int parsePositiveInt(const char* value, const char* name) { char* end = nullptr; - const long parsed = std::strtol(value, &end, 10); + 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"); } diff --git a/benchmarks/cpu/lfm2_parallel_linear_shared_mx.cpp b/benchmarks/cpu/lfm2_parallel_linear_shared_mx.cpp index eba69bdd7..8044aec7b 100644 --- a/benchmarks/cpu/lfm2_parallel_linear_shared_mx.cpp +++ b/benchmarks/cpu/lfm2_parallel_linear_shared_mx.cpp @@ -60,7 +60,7 @@ float deterministicValue(uint32_t& state) { int parsePositiveInt(const char* value, const char* name) { char* end = nullptr; - const long parsed = std::strtol(value, &end, 10); + 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"); } From 9914c4432766372af8721d67dcab5e29448a313d Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Thu, 27 Aug 2026 11:55:38 +0800 Subject: [PATCH 11/21] fix(preprocessor): honour the checkpoint's ignore_merges flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The LFM2.5 checkpoint sets model.ignore_merges, which keeps a token that is already a vocabulary entry intact instead of rebuilding it from the merge table. The shared BPE ignored the flag, and roughly 2% of vocabulary entries longer than two characters cannot be reconstructed by merges alone, so ordinary prose produced different ids than the checkpoint's own tokenizer: "Croatia" became C/roat/ia rather than one token, and so did words like congruence, PREFIX, and Türkiye. Read the flag and short-circuit whole vocabulary entries when it is set. Every other checkpoint in the tree reports ignore_merges false, so their tokenization is bit-identical. The existing pinned-oracle strings happen to contain no merge-unreachable word, which is why they passed while the ids were wrong. Add a case that does contain one; it fails without this fix. --- mllm/preprocessor/tokenizers/BPE.cpp | 7 +++++++ mllm/preprocessor/tokenizers/BPE.hpp | 5 +++++ tests/cpu/Lfm2TokenizerTest.cpp | 13 +++++++++++++ 3 files changed, 25 insertions(+) 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/Lfm2TokenizerTest.cpp b/tests/cpu/Lfm2TokenizerTest.cpp index 38df63bb1..6075b79e6 100644 --- a/tests/cpu/Lfm2TokenizerTest.cpp +++ b/tests/cpu/Lfm2TokenizerTest.cpp @@ -57,6 +57,19 @@ TEST(Lfm2TokenizerTest, MatchesPinnedCheckpointOracleWhenProvided) { EXPECT_EQ(tool_input.ptr()[index], tool_expected[index]); } + // "Croatia" is a vocabulary entry the merge table cannot rebuild, so it only + // survives as one id when the checkpoint's ignore_merges flag is honoured. + // Both other oracle strings above happen to avoid such words, which is why a + // merge-only BPE passed them while producing wrong ids for ordinary prose. + auto merge_unreachable = tokenizer.convertMessage({.prompt = "Croatia joined the European Union in 2013."}).at("sequence"); + const std::vector merge_unreachable_expected = {124894, 124899, 5922, 207, 116168, 8904, 278, + 4964, 6188, 296, 229, 523, 27, 22, + 124900, 207, 124899, 63514, 207, 124901}; + ASSERT_EQ(merge_unreachable.shape()[1], merge_unreachable_expected.size()); + for (size_t index = 0; index < merge_unreachable_expected.size(); ++index) { + EXPECT_EQ(merge_unreachable.ptr()[index], merge_unreachable_expected[index]); + } + const std::string multilingual = "你好 LFM2.5!"; const auto ordinary_tokens = tokenizer.tokenize(multilingual); const auto ordinary_ids = tokenizer.convert2Ids(ordinary_tokens); From 3c7b8ae8b9b67c83d845fe34e18b2d2b6271dd72 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Thu, 27 Aug 2026 11:55:38 +0800 Subject: [PATCH 12/21] fix(cpu): release the fused projection's prefill workspace CPUParallelLinearOp cached its KleidiAI LHS-pack scratch for every M and only ever grew it, so a prefill-sized buffer stayed resident for the rest of the process. CPULinearOp already avoids this by returning a throwaway buffer whenever M != 1; the fused operation did not carry that policy over when the shared-input path was extended to prefill. On the 2.6B product configuration this pins a prefill workspace in each of the 38 fused projections while decode needs about two kilobytes per operation. Take M and apply the same policy. The workspace is fully rewritten before any tile reads it, so this does not affect results. --- mllm/backends/cpu/ops/ParallelLinearOp.cpp | 12 +++++++----- mllm/backends/cpu/ops/ParallelLinearOp.hpp | 7 +++++-- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/mllm/backends/cpu/ops/ParallelLinearOp.cpp b/mllm/backends/cpu/ops/ParallelLinearOp.cpp index 3b8cc0c43..ac4f0fd36 100644 --- a/mllm/backends/cpu/ops/ParallelLinearOp.cpp +++ b/mllm/backends/cpu/ops/ParallelLinearOp.cpp @@ -36,11 +36,13 @@ void CPUParallelLinearOp::load(const ParameterFile::ptr_t& ploader) { } } -Tensor CPUParallelLinearOp::acquireKaiWorkspace(int32_t workspace_size) { - if (kai_workspace_.isNil() || kai_workspace_.numel() < static_cast(workspace_size)) { - kai_workspace_ = Tensor::empty({workspace_size}, kInt8, kCPU).alloc(); +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_workspace_; + return kai_decode_workspace_; } bool CPUParallelLinearOp::tryForwardSharedInputKai(const Tensor& input, std::vector& outputs) { @@ -93,7 +95,7 @@ bool CPUParallelLinearOp::tryForwardSharedInputKai(const Tensor& input, std::vec KaiHelper kai_helper; const size_t workspace_size = kai_helper.workspace_size(m, options_.in_channels, tile); if (workspace_size == 0 || workspace_size > static_cast(std::numeric_limits::max())) { return false; } - auto workspace = acquireKaiWorkspace(static_cast(workspace_size)); + auto workspace = acquireKaiWorkspace(static_cast(workspace_size), m); if (!kai_helper.matmul_shared_input(input.ptr(), projections.data(), weights_.size(), workspace.ptr(), m, options_.in_channels, tile, thread_count)) { return false; diff --git a/mllm/backends/cpu/ops/ParallelLinearOp.hpp b/mllm/backends/cpu/ops/ParallelLinearOp.hpp index 394611f1c..67a5754ae 100644 --- a/mllm/backends/cpu/ops/ParallelLinearOp.hpp +++ b/mllm/backends/cpu/ops/ParallelLinearOp.hpp @@ -20,10 +20,13 @@ class CPUParallelLinearOp final : public aops::ParallelLinearOp { private: bool tryForwardSharedInputKai(const Tensor& input, std::vector& outputs); - Tensor acquireKaiWorkspace(int32_t workspace_size); + Tensor acquireKaiWorkspace(int32_t workspace_size, int m); std::vector> fallback_ops_; - Tensor kai_workspace_; + // 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 { From 46dce91a4771f38444018d415d38a7cb438265fb Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Thu, 27 Aug 2026 11:55:38 +0800 Subject: [PATCH 13/21] fix(benchmark): initialize the runtime context before measuring The parallel-linear driver never called mllm::initializeContext(), so on the default threading vendor every tile-parallel call aborts and the driver only survives at threads=1 - the one setting its multi-worker screen is not about. Its shared-input sibling and the other CPU benchmarks already initialize the context. --- benchmarks/cpu/lfm2_parallel_linear.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/benchmarks/cpu/lfm2_parallel_linear.cpp b/benchmarks/cpu/lfm2_parallel_linear.cpp index 5320bd266..b864c9e11 100644 --- a/benchmarks/cpu/lfm2_parallel_linear.cpp +++ b/benchmarks/cpu/lfm2_parallel_linear.cpp @@ -16,6 +16,7 @@ #include #include "mllm/backends/cpu/kernels/arm/linear/kai.hpp" +#include "mllm/mllm.hpp" namespace { @@ -250,6 +251,7 @@ void runPair(std::string_view pair_name, std::string_view baseline, std::string_ 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; From 4010e04362b9dd6444001d659f0a316771f893cd Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Thu, 27 Aug 2026 11:55:39 +0800 Subject: [PATCH 14/21] fix(lfm2): reject a zero head count before dividing by it head_dim's default divides hidden_size by num_attention_heads while parsing. That expression is an ordinary function argument, so it is evaluated whether or not the config supplies head_dim, and it runs long before validate() can reject the value. A config with num_attention_heads set to zero therefore divided by zero instead of throwing. Check it where it is read. --- mllm/models/lfm2/configuration_lfm2.hpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mllm/models/lfm2/configuration_lfm2.hpp b/mllm/models/lfm2/configuration_lfm2.hpp index 383d85b23..067c2feca 100644 --- a/mllm/models/lfm2/configuration_lfm2.hpp +++ b/mllm/models/lfm2/configuration_lfm2.hpp @@ -23,6 +23,8 @@ struct Lfm2Config : protected ConfigFile { 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"); From 6320c44e0473b5c3a127975a000f8c0c0e4e169b Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Thu, 27 Aug 2026 15:40:17 +0800 Subject: [PATCH 15/21] test(cpu): align LFM2 coverage with upstream layout --- tests/cpu/CMakeLists.txt | 22 --- tests/cpu/CausalDepthwiseConvKernelTest.cpp | 94 ------------ tests/cpu/KernelTest.cpp | 82 ++++++++++ tests/cpu/Lfm2ConfigTest.cpp | 64 -------- tests/cpu/Lfm2RegisteredOpsTest.cpp | 160 -------------------- tests/cpu/Lfm2ShortConvTest.cpp | 73 --------- tests/cpu/Lfm2TokenizerTest.cpp | 82 ---------- 7 files changed, 82 insertions(+), 495 deletions(-) delete mode 100644 tests/cpu/CausalDepthwiseConvKernelTest.cpp delete mode 100644 tests/cpu/Lfm2ConfigTest.cpp delete mode 100644 tests/cpu/Lfm2RegisteredOpsTest.cpp delete mode 100644 tests/cpu/Lfm2ShortConvTest.cpp delete mode 100644 tests/cpu/Lfm2TokenizerTest.cpp diff --git a/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index b028c234d..90ce8037b 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -10,10 +10,6 @@ add_executable(Mllm-Test-Qwen35-GDN-Conv Qwen35GDNConvTest.cpp) target_link_libraries(Mllm-Test-Qwen35-GDN-Conv PRIVATE gtest_main MllmCPUBackend) target_include_directories(Mllm-Test-Qwen35-GDN-Conv PRIVATE ${MLLM_INCLUDE_DIR}) -add_executable(Mllm-Test-CausalDepthwiseConvKernel CausalDepthwiseConvKernelTest.cpp) -target_link_libraries(Mllm-Test-CausalDepthwiseConvKernel PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-CausalDepthwiseConvKernel PRIVATE ${MLLM_INCLUDE_DIR}) - add_executable(Mllm-Test-KaiW4A32Pack KaiW4A32PackTest.cpp) target_link_libraries(Mllm-Test-KaiW4A32Pack PRIVATE gtest_main MllmCPUBackend) target_include_directories(Mllm-Test-KaiW4A32Pack PRIVATE ${MLLM_INCLUDE_DIR}) @@ -44,24 +40,6 @@ target_include_directories(Mllm-Test-MiniCPM5-Model PRIVATE ${MLLM_INCLUDE_DIR}) target_compile_definitions(Mllm-Test-MiniCPM5-Model PRIVATE MINICPM5_EXAMPLE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../../examples/minicpm5") -add_executable(Mllm-Test-Lfm2-Config Lfm2ConfigTest.cpp) -target_link_libraries(Mllm-Test-Lfm2-Config PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-Lfm2-Config PRIVATE ${MLLM_INCLUDE_DIR}) -target_compile_definitions(Mllm-Test-Lfm2-Config - PRIVATE LFM2_EXAMPLE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../../examples/lfm2") - -add_executable(Mllm-Test-Lfm2-Tokenizer Lfm2TokenizerTest.cpp) -target_link_libraries(Mllm-Test-Lfm2-Tokenizer PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-Lfm2-Tokenizer PRIVATE ${MLLM_INCLUDE_DIR}) - -add_executable(Mllm-Test-Lfm2-ShortConv Lfm2ShortConvTest.cpp) -target_link_libraries(Mllm-Test-Lfm2-ShortConv PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-Lfm2-ShortConv PRIVATE ${MLLM_INCLUDE_DIR}) - -add_executable(Mllm-Test-Lfm2-RegisteredOps Lfm2RegisteredOpsTest.cpp) -target_link_libraries(Mllm-Test-Lfm2-RegisteredOps PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-Lfm2-RegisteredOps PRIVATE ${MLLM_INCLUDE_DIR}) - add_executable(Mllm-Test-CPUContiguousOp ContiguousOpTest.cpp) target_link_libraries(Mllm-Test-CPUContiguousOp PRIVATE gtest_main MllmRT MllmCPUBackend) target_include_directories(Mllm-Test-CPUContiguousOp PRIVATE ${MLLM_INCLUDE_DIR}) diff --git a/tests/cpu/CausalDepthwiseConvKernelTest.cpp b/tests/cpu/CausalDepthwiseConvKernelTest.cpp deleted file mode 100644 index e40823707..000000000 --- a/tests/cpu/CausalDepthwiseConvKernelTest.cpp +++ /dev/null @@ -1,94 +0,0 @@ -// Copyright (c) MLLM Team. -// Licensed under the MIT License. - -// Focused oracle for the history-first depthwise causal convolution kernel. -// -// The reference below is an independent scalar implementation of the frozen -// contract. It is deliberately not routed through the production kernel, so a -// vectorized fast path cannot validate itself. Both the output and the final -// history are compared bitwise: an output-only comparison would miss a -// corrupted history that only shows up in the next chunk. - -#include - -#include -#include - -#include "mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp" - -namespace { - -using mllm::cpu::causal_conv::depthwiseCausalConvHistoryFirstF32; - -// Deterministic index-derived fill. No RNG, so every host reproduces the same -// bytes without carrying a seed through the evidence record. -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; -} - -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; -} - -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]; - } - } - } -} - -TEST(CausalDepthwiseConvKernelTest, HistoryFirstK3MatchesScalarReferenceBitwise) { - constexpr int kKernel = 3; - for (int batch : {1, 2}) { - for (int sequence : {1, 2, 28, 225}) { - for (int channels : {1, 3, 4, 5, 2045, 2048}) { - for (bool non_zero_history : {false, true}) { - 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); - 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); - - ASSERT_EQ(kernel_output, reference_output) - << "history-first output mismatch for B=" << batch << " S=" << sequence << " C=" << channels - << " history=" << (non_zero_history ? "non-zero" : "zero"); - ASSERT_EQ(kernel_state, reference_state) - << "history-first state mismatch for B=" << batch << " S=" << sequence << " C=" << channels - << " history=" << (non_zero_history ? "non-zero" : "zero"); - } - } - } - } -} - -} // namespace diff --git a/tests/cpu/KernelTest.cpp b/tests/cpu/KernelTest.cpp index 9f8d613ee..d864fa667 100644 --- a/tests/cpu/KernelTest.cpp +++ b/tests/cpu/KernelTest.cpp @@ -3,12 +3,94 @@ #include +#include +#include + +#include "mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp" #include "mllm/mllm.hpp" #include "mllm/utils/CPUArchHelper.hpp" /// Kernel tests #include "ElementwiseKernelTest.hpp" +namespace { + +using mllm::cpu::causal_conv::depthwiseCausalConvHistoryFirstF32; + +float causalConvPatternValue(std::size_t index, int salt) { + const auto scaled = static_cast((index * 37U + static_cast(salt) * 11U) % 251U); + return (scaled - 125.0F) / 64.0F; +} + +std::vector makeCausalConvBuffer(std::size_t count, int salt) { + std::vector buffer(count); + for (std::size_t index = 0; index < count; ++index) { buffer[index] = causalConvPatternValue(index, salt); } + return buffer; +} + +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]; + } + } + } +} + +} // namespace + +//===----------------------------------------------------------------------===// +// Causal depthwise convolution +//===----------------------------------------------------------------------===// +TEST_F(KernelTest, CausalDepthwiseConvHistoryFirstK3MatchesScalarReferenceBitwise) { + constexpr int kKernel = 3; + for (int batch : {1, 2}) { + for (int sequence : {1, 2, 28, 225}) { + for (int channels : {1, 3, 4, 5, 2045, 2048}) { + for (bool non_zero_history : {false, true}) { + const auto element_count = static_cast(batch) * sequence * channels; + const auto state_count = static_cast(batch) * channels * (kKernel - 1); + const std::vector input = makeCausalConvBuffer(element_count, channels + sequence); + const std::vector weight = makeCausalConvBuffer(static_cast(channels) * kKernel, kKernel); + const std::vector initial_state = + non_zero_history ? makeCausalConvBuffer(state_count, 19) : std::vector(state_count, 0.0F); + + auto kernel_state = initial_state; + std::vector kernel_output(element_count, 0.0F); + 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); + + ASSERT_EQ(kernel_output, reference_output) + << "history-first output mismatch for B=" << batch << " S=" << sequence << " C=" << channels + << " history=" << (non_zero_history ? "non-zero" : "zero"); + ASSERT_EQ(kernel_state, reference_state) + << "history-first state mismatch for B=" << batch << " S=" << sequence << " C=" << channels + << " history=" << (non_zero_history ? "non-zero" : "zero"); + } + } + } + } +} + //===----------------------------------------------------------------------===// // Element wise ADD. // diff --git a/tests/cpu/Lfm2ConfigTest.cpp b/tests/cpu/Lfm2ConfigTest.cpp deleted file mode 100644 index 6cd23231a..000000000 --- a/tests/cpu/Lfm2ConfigTest.cpp +++ /dev/null @@ -1,64 +0,0 @@ -// Copyright (c) MLLM Team. -// Licensed under the MIT License. -#include - -#include - -#include "mllm/backends/cpu/ops/LinearOp.hpp" -#include "mllm/mllm.hpp" -#include "mllm/models/lfm2/configuration_lfm2.hpp" -#include "mllm/models/lfm2/modeling_lfm2.hpp" - -namespace { - -auto loadConfig() -> mllm::models::lfm2::Lfm2Config { - return mllm::models::lfm2::Lfm2Config(std::string(LFM2_EXAMPLE_DIR) + "/config_2.6B_w4a32_kai.json"); -} - -} // namespace - -TEST(Lfm2ConfigTest, Official26BContractUsesCompactAttentionSlots) { - const auto config = loadConfig(); - EXPECT_TRUE(mllm::models::lfm2::matchesOfficialRuntimeContract(config)); - EXPECT_EQ(config.numAttentionLayers(), 8); - EXPECT_EQ(config.numConvLayers(), 22); - EXPECT_EQ(config.attentionSlotForPhysicalLayer(2), 0); - EXPECT_EQ(config.attentionSlotForPhysicalLayer(27), 7); - EXPECT_THROW((void)config.attentionSlotForPhysicalLayer(0), std::invalid_argument); -} - -TEST(Lfm2ConfigTest, NativeKVCacheUsesEightHeadsPerLogicalSlot) { - mllm::initializeContext(); - const auto config = loadConfig(); - auto model = mllm::models::lfm2::Lfm2ForCausalLM(config); - EXPECT_EQ(model.kvCache().getLayerNums(), 8); - EXPECT_EQ(model.kvCache().kvHeads(), 8); - EXPECT_EQ(model.kvCache().headDim(), 64); - EXPECT_EQ(model.kvCache().maxCacheLength(), 2048); - EXPECT_NO_THROW(model.resetState()); - EXPECT_EQ(model.kvCache().getCurrentSeqCnt(0), 0); - model.kvCache().setCurrentSeqCnt(1); - auto sequence = mllm::Tensor::zeros({1, 1}, mllm::kInt64, mllm::kCPU); - EXPECT_THROW((void)model.forward({{"sequence", sequence}}, {}), std::invalid_argument); - mllm::shutdownContext(); -} - -TEST(Lfm2ConfigTest, RejectsPhysicalLayerScheduleDrift) { - auto config = loadConfig(); - config.layer_types[0] = "full_attention"; - EXPECT_FALSE(mllm::models::lfm2::matchesOfficialRuntimeContract(config)); -} - -TEST(Lfm2ConfigTest, KaiW4A32ThreadCapsSeparateDecodeAndPrefill) { - using mllm::cpu::detail::kaiW4A32ThreadCount; - EXPECT_EQ(kaiW4A32ThreadCount(1, 8, 4, 6), 4); - EXPECT_EQ(kaiW4A32ThreadCount(28, 8, 4, 6), 6); - EXPECT_EQ(kaiW4A32ThreadCount(1, 2, 4, 6), 2); - EXPECT_EQ(kaiW4A32ThreadCount(28, 4, 4, 6), 4); - EXPECT_EQ(kaiW4A32ThreadCount(1, 8, 0, 0), 8); - EXPECT_EQ(kaiW4A32ThreadCount(28, 8, 0, 0), 8); - EXPECT_EQ(kaiW4A32ThreadCount(1, 8, 4, 0), 4); - EXPECT_EQ(kaiW4A32ThreadCount(28, 8, 0, 6), 6); - EXPECT_EQ(kaiW4A32ThreadCount(1, 8, 12, 12), 8); - EXPECT_EQ(kaiW4A32ThreadCount(28, 8, 12, 12), 8); -} diff --git a/tests/cpu/Lfm2RegisteredOpsTest.cpp b/tests/cpu/Lfm2RegisteredOpsTest.cpp deleted file mode 100644 index 0d0cee21a..000000000 --- a/tests/cpu/Lfm2RegisteredOpsTest.cpp +++ /dev/null @@ -1,160 +0,0 @@ -// Copyright (c) MLLM Team. -// Licensed under the MIT License. - -#include -#include - -#include -#include -#include -#include -#include -#include - -#include "mllm/compile/ir/Trace.hpp" -#include "mllm/core/aops/ParallelLinearOp.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 Lfm2RegisteredOpsTest : public testing::Test { - protected: - static void SetUpTestSuite() { mllm::initializeContext(); } -}; - -class CausalConvTraceModule final : public mllm::nn::Module { - public: - CausalConvTraceModule() : Module("causal_conv_trace") { - conv_ = reg("conv", 4, 3, false, true, - 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}; - } - - private: - mllm::nn::CausalDepthwiseConv1D conv_; -}; - -class ParallelLinearModule final : public mllm::nn::Module { - public: - explicit ParallelLinearModule(std::string name) : Module(std::move(name)) { - projections_ = reg( - "pair", mllm::aops::ParallelLinearOpOptions{.in_channels = 2, - .out_channels = {2, 1}, - .projection_names = {"left", "right"}, - .bias = false, - .impl_type = mllm::aops::LinearImplTypes::kGGUF, - .kai_w4a32_decode_thread_cap = 4, - .kai_w4a32_prefill_thread_cap = 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 tensor = Tensor::empty(shape, mllm::kFloat32, mllm::kCPU).setMemType(mllm::kParamsNormal).setName(name).alloc(); - std::copy(values.begin(), values.end(), tensor.ptr()); - return tensor; -} - -TEST_F(Lfm2RegisteredOpsTest, CausalConvTracesAndSerializesStateSemantics) { - CausalConvTraceModule module; - auto ir_context = mllm::ir::trace(module, Tensor::empty({1, 2, 4}, mllm::kFloat32, mllm::kCPU), - Tensor::empty({1, 4, 2}, mllm::kFloat32, mllm::kCPU)); - auto op = findOp(ir_context->topLevelOp()); - ASSERT_NE(op, nullptr); - const auto options = mllm::jit::binary::dumpLinalgIROptions(op); - EXPECT_EQ(options.at("channels"), 4); - EXPECT_EQ(options.at("kernel_size"), 3); - 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); - EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kCausalDepthwiseConv1D); -} - -TEST_F(Lfm2RegisteredOpsTest, ParallelLinearOwnsSiblingParametersAndFallsBackCorrectly) { - ParallelLinearModule module("parallel_eager"); - auto parameters = mllm::ParameterFile::create(); - parameters->push("parallel_eager.left.weight", parameter("parallel_eager.left.weight", {2, 2}, {1.0F, 2.0F, 3.0F, 4.0F})); - parameters->push("parallel_eager.right.weight", parameter("parallel_eager.right.weight", {1, 2}, {5.0F, 6.0F})); - module.load(parameters); - - auto input = Tensor::empty({1, 1, 2}, mllm::kFloat32, mllm::kCPU).alloc(); - input.ptr()[0] = 2.0F; - input.ptr()[1] = 3.0F; - const auto outputs = module(input); - ASSERT_EQ(outputs.size(), 2); - EXPECT_EQ(outputs[0].shape(), (Tensor::shape_t{1, 1, 2})); - EXPECT_FLOAT_EQ(outputs[0].ptr()[0], 8.0F); - EXPECT_FLOAT_EQ(outputs[0].ptr()[1], 18.0F); - EXPECT_FLOAT_EQ(outputs[1].ptr()[0], 28.0F); -} - -TEST_F(Lfm2RegisteredOpsTest, ParallelLinearTracesAndSerializesProjectionContract) { - ParallelLinearModule module("parallel_trace"); - 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); - const auto options = mllm::jit::binary::dumpLinalgIROptions(op); - EXPECT_EQ(options.at("out_channels"), (std::vector{2, 1})); - EXPECT_EQ(options.at("projection_names"), (std::vector{"left", "right"})); - EXPECT_EQ(options.at("kai_w4a32_decode_thread_cap"), 4); - EXPECT_EQ(options.at("kai_w4a32_prefill_thread_cap"), 6); - - const auto restored = mllm::jit::interpreter::aopsFromJson( - nlohmann::json{{"op_type", "ParallelLinear"}, {"backend", "CPU"}, {"op_options", options}}); - ASSERT_NE(restored, nullptr); - EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kParallelLinear); -} - -// Fused parameters resolve in the parent scope, so ambiguous or scope-escaping -// projection names would silently bind the wrong checkpoint tensors. -TEST_F(Lfm2RegisteredOpsTest, ParallelLinearRejectsAmbiguousProjectionNames) { - auto reshapeWith = [](std::vector projection_names) { - auto op = std::make_shared( - mllm::aops::ParallelLinearOpOptions{.in_channels = 2, - .out_channels = {2, 1}, - .projection_names = std::move(projection_names), - .bias = false, - .impl_type = mllm::aops::LinearImplTypes::kGGUF}); - std::vector inputs = {Tensor::empty({1, 1, 2}, mllm::kFloat32, mllm::kCPU)}; - std::vector outputs; - op->reshape(inputs, outputs); - }; - - EXPECT_THROW(reshapeWith({"same", "same"}), std::invalid_argument); - EXPECT_THROW(reshapeWith({"left", "nested.right"}), std::invalid_argument); - EXPECT_THROW(reshapeWith({"left", ""}), std::invalid_argument); - EXPECT_NO_THROW(reshapeWith({"left", "right"})); -} - -} // namespace diff --git a/tests/cpu/Lfm2ShortConvTest.cpp b/tests/cpu/Lfm2ShortConvTest.cpp deleted file mode 100644 index 315c5c723..000000000 --- a/tests/cpu/Lfm2ShortConvTest.cpp +++ /dev/null @@ -1,73 +0,0 @@ -// Copyright (c) MLLM Team. -// Licensed under the MIT License. -#include - -#include -#include -#include - -#include "mllm/mllm.hpp" -#include "mllm/models/lfm2/modeling_lfm2.hpp" - -namespace { - -class Lfm2ShortConvTest : public testing::Test { - protected: - static void SetUpTestSuite() { mllm::initializeContext(); } -}; - -auto tensor(const std::string& name, const mllm::Tensor::shape_t& shape, const std::vector& values) -> mllm::Tensor { - auto result = mllm::Tensor::empty(shape, mllm::kFloat32, mllm::kCPU).setMemType(mllm::kParamsNormal).setName(name).alloc(); - EXPECT_EQ(result.numel(), values.size()); - std::copy(values.begin(), values.end(), result.ptr()); - return result; -} - -auto input(const std::vector& values) -> mllm::Tensor { - auto result = mllm::Tensor::empty({1, static_cast(values.size()), 1}, mllm::kFloat32, mllm::kCPU).alloc(); - std::copy(values.begin(), values.end(), result.ptr()); - return result; -} - -auto shortConv() -> mllm::models::lfm2::Lfm2ShortConv { - mllm::models::lfm2::Lfm2Config config; - config.hidden_size = 1; - config.conv_L_cache = 3; - config.conv_bias = false; - // Keep this semantics-only test portable. kDefault selects the ARM-only - // MllmBlas fallback for the deliberately tiny K=1 geometry on non-BLAS x86. - config.linear_impl_type = mllm::aops::LinearImplTypes::kGGUF; - auto module = mllm::models::lfm2::Lfm2ShortConv("unit", config); - auto parameters = mllm::ParameterFile::create(); - parameters->push("unit.in_proj.weight", tensor("unit.in_proj.weight", {3, 1}, {1.0F, 1.0F, 1.0F})); - parameters->push("unit.conv.weight", tensor("unit.conv.weight", {1, 1, 3}, {1.0F, 2.0F, 3.0F})); - parameters->push("unit.out_proj.weight", tensor("unit.out_proj.weight", {1, 1}, {1.0F})); - module.load(parameters); - return module; -} - -auto values(mllm::Tensor output) -> std::vector { - output = output.contiguous(); - return {output.ptr(), output.ptr() + output.numel()}; -} - -TEST_F(Lfm2ShortConvTest, ChunkedPrefillAndDecodeMatchOneShotCausalConvolution) { - auto chunked = shortConv(); - auto prefill = values(chunked(input({1.0F, 2.0F}))[0]); - auto decode = values(chunked(input({3.0F}))[0]); - EXPECT_EQ(prefill, (std::vector{3.0F, 28.0F})); - EXPECT_EQ(decode, (std::vector{108.0F})); - EXPECT_EQ(values(chunked.state()), (std::vector{4.0F, 9.0F})); - - auto one_shot = shortConv(); - EXPECT_EQ(values(one_shot(input({1.0F, 2.0F, 3.0F}))[0]), (std::vector{3.0F, 28.0F, 108.0F})); -} - -TEST_F(Lfm2ShortConvTest, ResetClearsTheTwoRequiredHistoricalSamples) { - auto module = shortConv(); - (void)module(input({2.0F})); - module.resetState(1); - EXPECT_EQ(values(module.state()), (std::vector{0.0F, 0.0F})); -} - -} // namespace diff --git a/tests/cpu/Lfm2TokenizerTest.cpp b/tests/cpu/Lfm2TokenizerTest.cpp deleted file mode 100644 index 6075b79e6..000000000 --- a/tests/cpu/Lfm2TokenizerTest.cpp +++ /dev/null @@ -1,82 +0,0 @@ -// Copyright (c) MLLM Team. -// Licensed under the MIT License. -#include - -#include -#include -#include - -#include "mllm/mllm.hpp" -#include "mllm/models/lfm2/tokenization_lfm2.hpp" - -TEST(Lfm2TokenizerTest, GroupsDigitsInRunsOfAtMostThree) { - std::vector pieces; - ASSERT_TRUE(mllm::models::lfm2::tokenizerRegex("1234567", pieces)); - EXPECT_EQ(pieces, (std::vector{L"123", L"456", L"7"})); -} - -TEST(Lfm2TokenizerTest, GenerationPromptEndsAtThinkingTokenWithoutNewline) { - const auto text = mllm::models::lfm2::Lfm2Message::render({.prompt = "Hello"}); - EXPECT_EQ(text.substr(text.size() - 7), ""); - EXPECT_EQ(text.find("<|startoftext|>"), 0); - EXPECT_EQ(text.find("<|im_start|>assistant\n"), text.size() - 29); -} - -TEST(Lfm2TokenizerTest, RendersPinnedSystemAndRawToolSchemaContract) { - const auto text = mllm::models::lfm2::Lfm2Message::render( - {.prompt = "Weather?", .system_prompt = "Be concise.", .tools = {R"({"type": "function"})"}}); - EXPECT_EQ(text, "<|startoftext|><|im_start|>system\nBe concise.\nList of tools: [{\"type\": \"function\"}]<|im_end|>\n" - "<|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n"); -} - -TEST(Lfm2TokenizerTest, StreamsUtf8AcrossTokenBoundaries) { - mllm::models::lfm2::StreamingUtf8Decoder decoder; - EXPECT_EQ(decoder.append("\xF0\x9F"), ""); - EXPECT_EQ(decoder.append("\x98\x80"), "\xF0\x9F\x98\x80"); - EXPECT_EQ(decoder.finish(), ""); -} - -TEST(Lfm2TokenizerTest, MatchesPinnedCheckpointOracleWhenProvided) { - const char* tokenizer_path = std::getenv("MLLM_LFM2_TOKENIZER_JSON"); - if (tokenizer_path == nullptr) GTEST_SKIP() << "set MLLM_LFM2_TOKENIZER_JSON to run checkpoint oracle"; - mllm::initializeContext(); - auto tokenizer = mllm::models::lfm2::Lfm2Tokenizer(tokenizer_path); - auto input = tokenizer.convertMessage({.prompt = "Hello"}).at("sequence"); - const std::vector expected = {124894, 124899, 5922, 207, 35808, 124900, 207, 124899, 63514, 207, 124901}; - ASSERT_EQ(input.shape()[1], expected.size()); - for (size_t index = 0; index < expected.size(); ++index) EXPECT_EQ(input.ptr()[index], expected[index]); - - auto tool_input = - tokenizer.convertMessage({.prompt = "Weather?", .system_prompt = "Be concise.", .tools = {R"({"type": "function"})"}}) - .at("sequence"); - const std::vector tool_expected = {124894, 124899, 23630, 207, 4184, 55911, 318, 3120, 302, 5985, - 34, 66155, 5882, 6380, 496, 5545, 66212, 124900, 207, 124899, - 5922, 207, 97056, 39, 124900, 207, 124899, 63514, 207, 124901}; - ASSERT_EQ(tool_input.shape()[1], tool_expected.size()); - for (size_t index = 0; index < tool_expected.size(); ++index) { - EXPECT_EQ(tool_input.ptr()[index], tool_expected[index]); - } - - // "Croatia" is a vocabulary entry the merge table cannot rebuild, so it only - // survives as one id when the checkpoint's ignore_merges flag is honoured. - // Both other oracle strings above happen to avoid such words, which is why a - // merge-only BPE passed them while producing wrong ids for ordinary prose. - auto merge_unreachable = tokenizer.convertMessage({.prompt = "Croatia joined the European Union in 2013."}).at("sequence"); - const std::vector merge_unreachable_expected = {124894, 124899, 5922, 207, 116168, 8904, 278, - 4964, 6188, 296, 229, 523, 27, 22, - 124900, 207, 124899, 63514, 207, 124901}; - ASSERT_EQ(merge_unreachable.shape()[1], merge_unreachable_expected.size()); - for (size_t index = 0; index < merge_unreachable_expected.size(); ++index) { - EXPECT_EQ(merge_unreachable.ptr()[index], merge_unreachable_expected[index]); - } - - const std::string multilingual = "你好 LFM2.5!"; - const auto ordinary_tokens = tokenizer.tokenize(multilingual); - const auto ordinary_ids = tokenizer.convert2Ids(ordinary_tokens); - std::string reconstructed; - for (int32_t index = 0; index < ordinary_ids.shape()[1]; ++index) { - reconstructed += tokenizer.detokenizeBytes(ordinary_ids.ptr()[index]); - } - EXPECT_EQ(reconstructed, multilingual); - mllm::shutdownContext(); -} From 23e88b15d6e7e924dcddbdb774e3375c0b4b302b Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Thu, 27 Aug 2026 15:52:00 +0800 Subject: [PATCH 16/21] test(cpu): keep the unified kernel registry concise --- tests/cpu/CausalDepthwiseConvKernelTest.hpp | 96 ++++++++++++++++++ tests/cpu/KernelTest.cpp | 104 +++++--------------- 2 files changed, 118 insertions(+), 82 deletions(-) create mode 100644 tests/cpu/CausalDepthwiseConvKernelTest.hpp diff --git a/tests/cpu/CausalDepthwiseConvKernelTest.hpp b/tests/cpu/CausalDepthwiseConvKernelTest.hpp new file mode 100644 index 000000000..2b3f2567a --- /dev/null +++ b/tests/cpu/CausalDepthwiseConvKernelTest.hpp @@ -0,0 +1,96 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +#include "mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp" + +#include "KernelTestHelper.hpp" + +class CausalDepthwiseConvKernelTest : public KernelTest { + public: + CausalDepthwiseConvKernelTest() = default; + ~CausalDepthwiseConvKernelTest() override = default; + + bool testHistoryFirstK3MatchesScalarReferenceBitwise(const std::vector>& cfgs) { + for (const auto& cfg : cfgs) { + if (!testOneCase(cfg)) { return false; } + } + return true; + } + + private: + static bool testOneCase(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 || kernel_state != reference_state) { + std::cerr << "history-first mismatch for B=" << batch << " S=" << sequence << " C=" << channels + << " history=" << (non_zero_history ? "non-zero" : "zero") << '\n'; + return false; + } + return true; + } + + 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 d864fa667..69285bb56 100644 --- a/tests/cpu/KernelTest.cpp +++ b/tests/cpu/KernelTest.cpp @@ -3,94 +3,12 @@ #include -#include -#include - -#include "mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp" #include "mllm/mllm.hpp" #include "mllm/utils/CPUArchHelper.hpp" /// Kernel tests #include "ElementwiseKernelTest.hpp" -namespace { - -using mllm::cpu::causal_conv::depthwiseCausalConvHistoryFirstF32; - -float causalConvPatternValue(std::size_t index, int salt) { - const auto scaled = static_cast((index * 37U + static_cast(salt) * 11U) % 251U); - return (scaled - 125.0F) / 64.0F; -} - -std::vector makeCausalConvBuffer(std::size_t count, int salt) { - std::vector buffer(count); - for (std::size_t index = 0; index < count; ++index) { buffer[index] = causalConvPatternValue(index, salt); } - return buffer; -} - -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]; - } - } - } -} - -} // namespace - -//===----------------------------------------------------------------------===// -// Causal depthwise convolution -//===----------------------------------------------------------------------===// -TEST_F(KernelTest, CausalDepthwiseConvHistoryFirstK3MatchesScalarReferenceBitwise) { - constexpr int kKernel = 3; - for (int batch : {1, 2}) { - for (int sequence : {1, 2, 28, 225}) { - for (int channels : {1, 3, 4, 5, 2045, 2048}) { - for (bool non_zero_history : {false, true}) { - const auto element_count = static_cast(batch) * sequence * channels; - const auto state_count = static_cast(batch) * channels * (kKernel - 1); - const std::vector input = makeCausalConvBuffer(element_count, channels + sequence); - const std::vector weight = makeCausalConvBuffer(static_cast(channels) * kKernel, kKernel); - const std::vector initial_state = - non_zero_history ? makeCausalConvBuffer(state_count, 19) : std::vector(state_count, 0.0F); - - auto kernel_state = initial_state; - std::vector kernel_output(element_count, 0.0F); - 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); - - ASSERT_EQ(kernel_output, reference_output) - << "history-first output mismatch for B=" << batch << " S=" << sequence << " C=" << channels - << " history=" << (non_zero_history ? "non-zero" : "zero"); - ASSERT_EQ(kernel_state, reference_state) - << "history-first state mismatch for B=" << batch << " S=" << sequence << " C=" << channels - << " history=" << (non_zero_history ? "non-zero" : "zero"); - } - } - } - } -} - //===----------------------------------------------------------------------===// // Element wise ADD. // @@ -615,6 +533,28 @@ TEST_F(ElementwiseKernelTest, DivScalarInt32) { true); } +//===----------------------------------------------------------------------===// +// Causal depthwise convolution +//===----------------------------------------------------------------------===// +#include "CausalDepthwiseConvKernelTest.hpp" +TEST_F(CausalDepthwiseConvKernelTest, HistoryFirstK3MatchesScalarReferenceBitwise) { + EXPECT_EQ(testHistoryFirstK3MatchesScalarReferenceBitwise({ + {{"B", 1}, {"S", 1}, {"C", 1}, {"non_zero_history", 0}}, + {{"B", 1}, {"S", 1}, {"C", 1}, {"non_zero_history", 1}}, + {{"B", 2}, {"S", 2}, {"C", 3}, {"non_zero_history", 0}}, + {{"B", 2}, {"S", 2}, {"C", 3}, {"non_zero_history", 1}}, + {{"B", 1}, {"S", 28}, {"C", 4}, {"non_zero_history", 0}}, + {{"B", 1}, {"S", 28}, {"C", 4}, {"non_zero_history", 1}}, + {{"B", 2}, {"S", 2}, {"C", 5}, {"non_zero_history", 0}}, + {{"B", 2}, {"S", 2}, {"C", 5}, {"non_zero_history", 1}}, + {{"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); +} + //===----------------------------------------------------------------------===// // CausalMaskOp //===----------------------------------------------------------------------===// From 6b6b1965fda4925b0f84cee2efcea84c52005c26 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Thu, 27 Aug 2026 18:41:11 +0800 Subject: [PATCH 17/21] test(nn): cover causal conv and parallel linear ops --- tests/cpu/CausalDepthwiseConvKernelTest.hpp | 31 +-- tests/cpu/KernelTest.cpp | 13 +- tests/nn/CMakeLists.txt | 8 + tests/nn/CausalDepthwiseConv1DTest.cpp | 211 ++++++++++++++++++++ tests/nn/ParallelLinearTest.cpp | 147 ++++++++++++++ 5 files changed, 394 insertions(+), 16 deletions(-) create mode 100644 tests/nn/CausalDepthwiseConv1DTest.cpp create mode 100644 tests/nn/ParallelLinearTest.cpp diff --git a/tests/cpu/CausalDepthwiseConvKernelTest.hpp b/tests/cpu/CausalDepthwiseConvKernelTest.hpp index 2b3f2567a..2c81b52e9 100644 --- a/tests/cpu/CausalDepthwiseConvKernelTest.hpp +++ b/tests/cpu/CausalDepthwiseConvKernelTest.hpp @@ -6,12 +6,12 @@ #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" @@ -20,15 +20,7 @@ class CausalDepthwiseConvKernelTest : public KernelTest { CausalDepthwiseConvKernelTest() = default; ~CausalDepthwiseConvKernelTest() override = default; - bool testHistoryFirstK3MatchesScalarReferenceBitwise(const std::vector>& cfgs) { - for (const auto& cfg : cfgs) { - if (!testOneCase(cfg)) { return false; } - } - return true; - } - - private: - static bool testOneCase(const std::unordered_map& cfg) { + bool testHistoryFirstK3Once(const std::unordered_map& cfg) { constexpr int kKernel = 3; const int batch = cfg.at("B"); const int sequence = cfg.at("S"); @@ -52,14 +44,27 @@ class CausalDepthwiseConvKernelTest : public KernelTest { referenceDepthwiseCausalConvHistoryFirst(input, weight, reference_state, reference_output, batch, sequence, channels, kKernel); - if (kernel_output != reference_output || kernel_state != reference_state) { - std::cerr << "history-first mismatch for B=" << batch << " S=" << sequence << " C=" << channels - << " history=" << (non_zero_history ? "non-zero" : "zero") << '\n'; + 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; diff --git a/tests/cpu/KernelTest.cpp b/tests/cpu/KernelTest.cpp index 69285bb56..e01bfabe2 100644 --- a/tests/cpu/KernelTest.cpp +++ b/tests/cpu/KernelTest.cpp @@ -538,15 +538,22 @@ TEST_F(ElementwiseKernelTest, DivScalarInt32) { //===----------------------------------------------------------------------===// #include "CausalDepthwiseConvKernelTest.hpp" TEST_F(CausalDepthwiseConvKernelTest, HistoryFirstK3MatchesScalarReferenceBitwise) { - EXPECT_EQ(testHistoryFirstK3MatchesScalarReferenceBitwise({ + 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}}, - {{"B", 2}, {"S", 2}, {"C", 3}, {"non_zero_history", 0}}, - {{"B", 2}, {"S", 2}, {"C", 3}, {"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}}, 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/ParallelLinearTest.cpp b/tests/nn/ParallelLinearTest.cpp new file mode 100644 index 000000000..e73655d17 --- /dev/null +++ b/tests/nn/ParallelLinearTest.cpp @@ -0,0 +1,147 @@ +// 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(); } +}; + +mllm::aops::ParallelLinearOpOptions options(bool bias) { + return {.in_channels = 2, + .out_channels = {2, 1}, + .projection_names = {"left", "right"}, + .bias = bias, + .impl_type = mllm::aops::LinearImplTypes::kGGUF, + .kai_w4a32_decode_thread_cap = 4, + .kai_w4a32_prefill_thread_cap = 6}; +} + +class ParallelLinearModule final : public mllm::nn::Module { + public: + ParallelLinearModule(std::string name, bool bias) : Module(std::move(name)) { + projections_ = reg("pair", options(bias)); + } + + 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 From b557aeee845ec7e1e573d680af4214a6c8cb1c4b Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Thu, 27 Aug 2026 21:57:24 +0800 Subject: [PATCH 18/21] refactor(cpu): decouple parallel linear dispatch --- .../kernels/arm/linear/parallel_linear.cpp | 88 +++++++++++++++++++ .../common/linear/kai_w4a32_dispatch.cpp | 40 +++++++++ .../common/linear/kai_w4a32_dispatch.hpp | 23 +++++ .../common/parallel_linear/shared_input.hpp | 37 ++++++++ .../cpu/kernels/x86/parallel_linear.cpp | 30 +++++++ mllm/backends/cpu/ops/LinearOp.cpp | 33 ++----- mllm/backends/cpu/ops/LinearOp.hpp | 17 ---- mllm/backends/cpu/ops/ParallelLinearOp.cpp | 58 +++--------- mllm/core/aops/ParallelLinearOp.hpp | 2 +- mllm/models/lfm2/modeling_lfm2.hpp | 29 +++--- mllm/nn/Layer.hpp | 7 ++ mllm/nn/layers/ParallelLinear.cpp | 15 +++- mllm/nn/layers/ParallelLinear.hpp | 11 ++- tests/cpu/KernelTest.cpp | 6 ++ tests/cpu/ParallelLinearKernelTest.hpp | 37 ++++++++ tests/nn/ParallelLinearTest.cpp | 14 +-- 16 files changed, 320 insertions(+), 127 deletions(-) create mode 100644 mllm/backends/cpu/kernels/arm/linear/parallel_linear.cpp create mode 100644 mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.cpp create mode 100644 mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.hpp create mode 100644 mllm/backends/cpu/kernels/common/parallel_linear/shared_input.hpp create mode 100644 mllm/backends/cpu/kernels/x86/parallel_linear.cpp create mode 100644 tests/cpu/ParallelLinearKernelTest.hpp 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/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/LinearOp.cpp b/mllm/backends/cpu/ops/LinearOp.cpp index df0db22f8..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,18 +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) { - if (detail::shouldUseKaiW4A32I8mmPrefill(m)) { return kKaiW4A32I8mmTile; } + if (kai_w4a32::shouldUseI8mmPrefill(m)) { return kKaiW4A32I8mmTile; } return kKaiW4A32DotProdTile; } @@ -60,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); } } @@ -69,21 +56,11 @@ void traceKaiW4A32PrefillTile(KaiW4A32Tile tile, int m, int k, int n, int thread } // namespace -bool detail::shouldUseKaiW4A32I8mmPrefill(int m) { -#if defined(MLLM_HOST_ARCH_ARM64) || defined(MLLM_HOST_ARCH_ARM) - static const bool disabled = environmentFlagEnabled("MLLM_KAI_PREFILL_I8MM_DISABLE"); - return shouldUseKaiW4A32I8mmPrefill(m, disabled, cpuSupportsI8mm()); -#else - (void)m; - return false; -#endif -} - CPULinearOp::CPULinearOp(const aops::LinearOpOptions& options) : LinearOp(options) {} int CPULinearOp::kaiW4A32ThreadCount(int m) const { - return detail::kaiW4A32ThreadCount(m, options_.getThreads(), options_.kai_w4a32_decode_thread_cap, - options_.kai_w4a32_prefill_thread_cap); + 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) { diff --git a/mllm/backends/cpu/ops/LinearOp.hpp b/mllm/backends/cpu/ops/LinearOp.hpp index 46e4df87e..1519772c8 100644 --- a/mllm/backends/cpu/ops/LinearOp.hpp +++ b/mllm/backends/cpu/ops/LinearOp.hpp @@ -7,23 +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; -} - -bool shouldUseKaiW4A32I8mmPrefill(int m); - -constexpr int kaiW4A32ThreadCount(int m, int requested_threads, int decode_thread_cap, int prefill_thread_cap) { - // Every dynamic-input W4A32 KAI tile interprets the optional caps through - // this helper; zero keeps the repository-wide requested thread count. - const int cap = m == 1 ? decode_thread_cap : prefill_thread_cap; - return cap > 0 && cap < requested_threads ? cap : requested_threads; -} - -} // namespace detail - class CPULinearOp final : public aops::LinearOp { public: explicit CPULinearOp(const aops::LinearOpOptions& options); diff --git a/mllm/backends/cpu/ops/ParallelLinearOp.cpp b/mllm/backends/cpu/ops/ParallelLinearOp.cpp index ac4f0fd36..2330a443a 100644 --- a/mllm/backends/cpu/ops/ParallelLinearOp.cpp +++ b/mllm/backends/cpu/ops/ParallelLinearOp.cpp @@ -3,14 +3,11 @@ #include "mllm/backends/cpu/ops/ParallelLinearOp.hpp" -#include -#include #include -#include -#include +#include #include -#include "mllm/backends/cpu/kernels/Kernels.hpp" +#include "mllm/backends/cpu/kernels/common/parallel_linear/shared_input.hpp" namespace mllm::cpu { @@ -46,12 +43,8 @@ Tensor CPUParallelLinearOp::acquireKaiWorkspace(int32_t workspace_size, int m) { } bool CPUParallelLinearOp::tryForwardSharedInputKai(const Tensor& input, std::vector& outputs) { -#if defined(MLLM_HOST_ARCH_ARM64) || defined(MLLM_HOST_ARCH_ARM) constexpr size_t kMaximumSharedProjections = 3; constexpr auto kRequiredImpl = aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32; - using KaiHelper = ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk; - constexpr auto kDecodeTile = KaiHelper::Tiles::qai8dxp1x8_qsi4c32p8x8_1x8x32; - constexpr auto kPrefillTile = KaiHelper::Tiles::qai8dxp4x8_qsi4c32p8x8_4x8x32; 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 @@ -63,12 +56,6 @@ bool CPUParallelLinearOp::tryForwardSharedInputKai(const Tensor& input, std::vec } const int32_t m = input.size(-2); if (m <= 0) { return false; } - const auto tile = m == 1 ? kDecodeTile : kPrefillTile; - if (m > 1 && !detail::shouldUseKaiW4A32I8mmPrefill(m)) { - // The generic shared-input dot-product prefill path has not passed the - // mobile product screen; preserve the established per-projection fallback. - 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() @@ -81,9 +68,12 @@ bool CPUParallelLinearOp::tryForwardSharedInputKai(const Tensor& input, std::vec } } - const int32_t thread_count = detail::kaiW4A32ThreadCount(m, options_.getThreads(), options_.kai_w4a32_decode_thread_cap, - options_.kai_w4a32_prefill_thread_cap); - std::array projections{}; + 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(), @@ -92,35 +82,9 @@ bool CPUParallelLinearOp::tryForwardSharedInputKai(const Tensor& input, std::vec }; } - KaiHelper kai_helper; - const size_t workspace_size = kai_helper.workspace_size(m, options_.in_channels, tile); - if (workspace_size == 0 || workspace_size > static_cast(std::numeric_limits::max())) { return false; } - auto workspace = acquireKaiWorkspace(static_cast(workspace_size), m); - if (!kai_helper.matmul_shared_input(input.ptr(), projections.data(), weights_.size(), workspace.ptr(), m, - options_.in_channels, tile, thread_count)) { - return false; - } - - static const bool trace_activation = [] { - const char* value = std::getenv("MLLM_KAI_SHARED_INPUT_TRACE"); - return value != nullptr && value[0] == '1' && value[1] == '\0'; - }(); - if (trace_activation) { - const uint32_t projection_group = static_cast(weights_.size() - 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", weights_.size(), m, - options_.in_channels, thread_count, m == 1 ? "dotprod_1x8" : "i8mm_4x8"); - } - } - return true; -#else - (void)input; - (void)outputs; - return false; -#endif + 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) { diff --git a/mllm/core/aops/ParallelLinearOp.hpp b/mllm/core/aops/ParallelLinearOp.hpp index 85a263908..7d7243743 100644 --- a/mllm/core/aops/ParallelLinearOp.hpp +++ b/mllm/core/aops/ParallelLinearOp.hpp @@ -34,7 +34,7 @@ class ParallelLinearOp : public BaseOp { void setup(const std::vector& inputs, std::vector& outputs) override; ParameterFile::ptr_t getParams() override; - inline const ParallelLinearOpOptions& options() const { return options_; } + [[nodiscard]] inline const ParallelLinearOpOptions& options() const { return options_; } protected: [[nodiscard]] std::string projectionParameterName(size_t index, const char* suffix) const; diff --git a/mllm/models/lfm2/modeling_lfm2.hpp b/mllm/models/lfm2/modeling_lfm2.hpp index 170472375..98e530106 100644 --- a/mllm/models/lfm2/modeling_lfm2.hpp +++ b/mllm/models/lfm2/modeling_lfm2.hpp @@ -24,6 +24,9 @@ namespace mllm::models::lfm2 { 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, @@ -32,20 +35,8 @@ inline auto makeLfm2LinearOptions(int32_t in_channels, int32_t out_channels, boo .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 = 4, - .kai_w4a32_prefill_thread_cap = 6}; -} - -inline auto makeLfm2ParallelLinearOptions(int32_t in_channels, std::vector out_channels, - std::vector projection_names, bool bias, aops::LinearImplTypes impl_type) - -> aops::ParallelLinearOpOptions { - return {.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 = 4, - .kai_w4a32_prefill_thread_cap = 6}; + .kai_w4a32_decode_thread_cap = kLfm2KaiDecodeThreadCap, + .kai_w4a32_prefill_thread_cap = kLfm2KaiPrefillThreadCap}; } class Lfm2MLP final : public nn::Module { @@ -53,8 +44,8 @@ class Lfm2MLP final : public nn::Module { Lfm2MLP() = default; Lfm2MLP(const std::string& name, const Lfm2Config& cfg) : nn::Module(name) { gate_up_proj_ = reg( - "gate_up_proj", makeLfm2ParallelLinearOptions(cfg.hidden_size, {cfg.intermediate_size, cfg.intermediate_size}, - {"w1", "w3"}, false, cfg.linear_impl_type)); + "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"); } @@ -78,9 +69,9 @@ class Lfm2Attention final : public nn::Module { query_heads_ = cfg.num_attention_heads; kv_heads_ = cfg.num_key_value_heads; qkv_proj_ = reg( - "qkv_proj", - makeLfm2ParallelLinearOptions(hidden_size_, {query_heads_ * head_dim_, kv_heads_ * head_dim_, kv_heads_ * head_dim_}, - {"q_proj", "k_proj", "v_proj"}, false, cfg.linear_impl_type)); + "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); 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/layers/ParallelLinear.cpp b/mllm/nn/layers/ParallelLinear.cpp index 36c739f27..2e6c015b9 100644 --- a/mllm/nn/layers/ParallelLinear.cpp +++ b/mllm/nn/layers/ParallelLinear.cpp @@ -3,10 +3,23 @@ #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(const aops::ParallelLinearOpOptions& options) : Layer(OpTypes::kParallelLinear, options) {} +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 index 11444080d..a8e56bae7 100644 --- a/mllm/nn/layers/ParallelLinear.hpp +++ b/mllm/nn/layers/ParallelLinear.hpp @@ -3,9 +3,11 @@ #pragma once +#include +#include #include -#include "mllm/core/aops/ParallelLinearOp.hpp" +#include "mllm/core/aops/LinearOp.hpp" #include "mllm/nn/Layer.hpp" namespace mllm::nn { @@ -13,9 +15,12 @@ namespace mllm::nn { class ParallelLinear : public Layer { public: ParallelLinear(); - explicit ParallelLinear(const aops::ParallelLinearOpOptions& options); - std::vector operator()(const Tensor& input) { return __main({input}); } + 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/tests/cpu/KernelTest.cpp b/tests/cpu/KernelTest.cpp index e01bfabe2..277bb51d7 100644 --- a/tests/cpu/KernelTest.cpp +++ b/tests/cpu/KernelTest.cpp @@ -562,6 +562,12 @@ TEST_F(CausalDepthwiseConvKernelTest, HistoryFirstK3MatchesScalarReferenceBitwis 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/ParallelLinearTest.cpp b/tests/nn/ParallelLinearTest.cpp index e73655d17..5a477b594 100644 --- a/tests/nn/ParallelLinearTest.cpp +++ b/tests/nn/ParallelLinearTest.cpp @@ -29,20 +29,12 @@ class ParallelLinearTest : public testing::Test { static void SetUpTestSuite() { mllm::initializeContext(); } }; -mllm::aops::ParallelLinearOpOptions options(bool bias) { - return {.in_channels = 2, - .out_channels = {2, 1}, - .projection_names = {"left", "right"}, - .bias = bias, - .impl_type = mllm::aops::LinearImplTypes::kGGUF, - .kai_w4a32_decode_thread_cap = 4, - .kai_w4a32_prefill_thread_cap = 6}; -} - class ParallelLinearModule final : public mllm::nn::Module { public: ParallelLinearModule(std::string name, bool bias) : Module(std::move(name)) { - projections_ = reg("pair", options(bias)); + 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 { From 65380a15fc74bb9a8d209a428c3a25a88cd0d812 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Thu, 27 Aug 2026 23:30:35 +0800 Subject: [PATCH 19/21] refactor(nn): encapsulate grouped attention options --- mllm/models/lfm2/modeling_lfm2.hpp | 4 +--- mllm/models/minicpm5/modeling_minicpm5.hpp | 4 +--- mllm/nn/layers/GroupedQueryAttention.cpp | 4 ++-- mllm/nn/layers/GroupedQueryAttention.hpp | 2 +- 4 files changed, 5 insertions(+), 9 deletions(-) diff --git a/mllm/models/lfm2/modeling_lfm2.hpp b/mllm/models/lfm2/modeling_lfm2.hpp index 98e530106..4c0f650e3 100644 --- a/mllm/models/lfm2/modeling_lfm2.hpp +++ b/mllm/models/lfm2/modeling_lfm2.hpp @@ -78,9 +78,7 @@ class Lfm2Attention final : public nn::Module { 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::GroupedQueryAttentionOpOptions{.implementation = aops::GroupedQueryAttentionImplementation::kDirectStrided}); + gqa_ = reg("gqa", aops::GroupedQueryAttentionImplementation::kDirectStrided); } std::vector forward(const std::vector& inputs, const std::vector& args) override { diff --git a/mllm/models/minicpm5/modeling_minicpm5.hpp b/mllm/models/minicpm5/modeling_minicpm5.hpp index 256f8d10a..bf976c4be 100644 --- a/mllm/models/minicpm5/modeling_minicpm5.hpp +++ b/mllm/models/minicpm5/modeling_minicpm5.hpp @@ -63,9 +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", - aops::GroupedQueryAttentionOpOptions{.implementation = aops::GroupedQueryAttentionImplementation::kDecodeNativeKV}); + gqa_decode_ = reg("gqa_decode", aops::GroupedQueryAttentionImplementation::kDecodeNativeKV); } std::vector forward(const std::vector& inputs, const std::vector& args) override { diff --git a/mllm/nn/layers/GroupedQueryAttention.cpp b/mllm/nn/layers/GroupedQueryAttention.cpp index bce4477a3..0e629c091 100644 --- a/mllm/nn/layers/GroupedQueryAttention.cpp +++ b/mllm/nn/layers/GroupedQueryAttention.cpp @@ -8,7 +8,7 @@ namespace mllm::nn { GroupedQueryAttention::GroupedQueryAttention() : Layer(OpTypes::kGroupedQueryAttention, aops::GroupedQueryAttentionOpOptions{}) {} -GroupedQueryAttention::GroupedQueryAttention(const aops::GroupedQueryAttentionOpOptions& options) - : Layer(OpTypes::kGroupedQueryAttention, options) {} +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 index 67f00bc37..99eca8769 100644 --- a/mllm/nn/layers/GroupedQueryAttention.hpp +++ b/mllm/nn/layers/GroupedQueryAttention.hpp @@ -11,7 +11,7 @@ namespace mllm::nn { class GroupedQueryAttention : public Layer { public: GroupedQueryAttention(); - explicit GroupedQueryAttention(const aops::GroupedQueryAttentionOpOptions& options); + explicit GroupedQueryAttention(aops::GroupedQueryAttentionImplementation implementation); MLLM_LAYER_ANY_INPUTS_1_OUTPUTS_FORWARD }; From 372313f56c3ef0ffdb4137c4386e4dfab047ab7b Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Thu, 27 Aug 2026 23:40:03 +0800 Subject: [PATCH 20/21] refactor(nn): hide causal convolution op options --- mllm/nn/layers/CausalDepthwiseConv1D.cpp | 3 --- mllm/nn/layers/CausalDepthwiseConv1D.hpp | 1 - 2 files changed, 4 deletions(-) diff --git a/mllm/nn/layers/CausalDepthwiseConv1D.cpp b/mllm/nn/layers/CausalDepthwiseConv1D.cpp index 6039c5867..461dd00ac 100644 --- a/mllm/nn/layers/CausalDepthwiseConv1D.cpp +++ b/mllm/nn/layers/CausalDepthwiseConv1D.cpp @@ -8,9 +8,6 @@ namespace mllm::nn { CausalDepthwiseConv1D::CausalDepthwiseConv1D() : Layer(OpTypes::kCausalDepthwiseConv1D, aops::CausalDepthwiseConv1DOpOptions{}) {} -CausalDepthwiseConv1D::CausalDepthwiseConv1D(const aops::CausalDepthwiseConv1DOpOptions& options) - : Layer(OpTypes::kCausalDepthwiseConv1D, options) {} - 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, diff --git a/mllm/nn/layers/CausalDepthwiseConv1D.hpp b/mllm/nn/layers/CausalDepthwiseConv1D.hpp index f6c79fff1..ddd76c143 100644 --- a/mllm/nn/layers/CausalDepthwiseConv1D.hpp +++ b/mllm/nn/layers/CausalDepthwiseConv1D.hpp @@ -11,7 +11,6 @@ namespace mllm::nn { class CausalDepthwiseConv1D : public Layer { public: CausalDepthwiseConv1D(); - explicit CausalDepthwiseConv1D(const aops::CausalDepthwiseConv1DOpOptions& options); CausalDepthwiseConv1D(int32_t channels, int32_t kernel_size, bool bias, bool state_inplace, aops::CausalDepthwiseConv1DAccumulationOrder accumulation_order); From 17897b99bf1db963d4a6ea605e7700cb663482fa Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Fri, 28 Aug 2026 01:19:44 +0800 Subject: [PATCH 21/21] refactor(cpu): scope ARM OpenMP to owning sources --- CMakeLists.txt | 8 ----- examples/lfm2/README.md | 12 ++++---- mllm/backends/cpu/CMakeLists.txt | 50 +++++++++++++++++++++++--------- 3 files changed, 43 insertions(+), 27 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 1d7d2f1f2..b16ae832a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -41,14 +41,6 @@ option(MLLM_BLAS_VENDOR_BLIS "Enable BLIS BLAS for multi-platform" OFF) # CPU Backend: SME2 and SVE2 option(MLLM_CPU_BACKEND_USE_SME2 "Enable SME2" OFF) -option( - MLLM_ARM_CPU_BACKEND_USE_OPENMP - "Compile ARM CPU backend operators and kernels with OpenMP" - ON) -option( - MLLM_ARM_KAI_USE_OPENMP - "Compile only the ARM KleidiAI linear-kernel translation unit with OpenMP" - ON) # Ascend Backend: Options option(MLLM_ASCEND_CPU_DEBUG_MODE "Enable CPU Debug mode in ascend" OFF) diff --git a/examples/lfm2/README.md b/examples/lfm2/README.md index f72128cee..f1fa3cc0c 100644 --- a/examples/lfm2/README.md +++ b/examples/lfm2/README.md @@ -57,11 +57,9 @@ build/bin/mllm-lfm2-runner \ --print_token_ids ``` -For the Android ARM build, keep runtime OpenMP enabled but configure the CPU -backend without backend-wide OpenMP. LFM2.5 still selects the shared W4A32 -I8MM prefill path and retained KAI decode workspace; avoiding OpenMP regions in -every backend operator preserves decode latency for its 8-attention / 22-conv -hybrid schedule. +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 \ @@ -71,7 +69,9 @@ cmake -S . -B build-android \ -DMLLM_CROSS_COMPILE=ON \ -DMLLM_BUILD_ARM_BACKEND=ON \ -DMLLM_ENABLE_EXAMPLE=ON \ - -DMLLM_ARM_CPU_BACKEND_USE_OPENMP=OFF + -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 ``` diff --git a/mllm/backends/cpu/CMakeLists.txt b/mllm/backends/cpu/CMakeLists.txt index 24c0a44c8..86623e2b3 100644 --- a/mllm/backends/cpu/CMakeLists.txt +++ b/mllm/backends/cpu/CMakeLists.txt @@ -147,20 +147,44 @@ 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 AND MLLM_ARM_CPU_BACKEND_USE_OPENMP) - 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}) - elseif(MLLM_BUILD_ARM_BACKEND AND MLLM_ARM_KAI_USE_OPENMP) - # The KAI linear helper owns its tile-parallel loops. Compiling only - # this translation unit with OpenMP lets quantized Linear reuse those - # loops without adding parallel-region overhead to every CPU operator. - set_source_files_properties( + 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 - PROPERTIES COMPILE_OPTIONS "${OpenMP_CXX_FLAGS}") + ${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_include_directories(MllmCPUBackend PUBLIC ${OpenMP_CXX_INCLUDE_DIR}) endif()