diff --git a/README-ZH.md b/README-ZH.md index 9106c4153..2bf6db4ae 100644 --- a/README-ZH.md +++ b/README-ZH.md @@ -107,6 +107,7 @@ mllm 框架可以与主流社区框架的模型检查点无缝集成。通过 ml | [Qwen3-4B](https://github.com/QwenLM/Qwen3) | [✔️ w4a8](https://www.modelscope.cn/models/mllmTeam/Qwen3-4B-w4a8-i8mm-kai) | | | | Qwen3.5-2B | | ✔️ W4A16 / W8A8 | | | Qwen3.5-4B | | ✔️ W4A16 / W8A8 | | +| [Ling-3.0-tiny](https://huggingface.co/inclusionAI/Ling-3.0-tiny) | [✔️ W4A32，ARM64 macOS/Android](./examples/ling3/README.md) | | | | Qwen3-VL-2B-Instruct | | ✔️ W4A16 / W8A8 | | | Qwen3-VL-4B-Instruct | | ✔️ W4A16 / W8A8 | | | [DeepSeek-OCR](https://github.com/deepseek-ai/DeepSeek-OCR) | [✔️ w4a8](https://www.modelscope.cn/models/mllmTeam/DeepSeek-OCR-w4a8-i8mm-kai) | | | diff --git a/README.md b/README.md index a511778c8..793b1b514 100644 --- a/README.md +++ b/README.md @@ -108,6 +108,7 @@ The mllm framework integrates seamlessly with popular community frameworks' chec | [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) | | | +| [Ling-3.0-tiny](https://huggingface.co/inclusionAI/Ling-3.0-tiny) | [✔️ W4A32, ARM64 macOS/Android](./examples/ling3/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/bench_assets/ling3_tiny_architecture.png b/bench_assets/ling3_tiny_architecture.png new file mode 100644 index 000000000..ef8fdeab3 Binary files /dev/null and b/bench_assets/ling3_tiny_architecture.png differ diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index ed1cbeb82..34e544b6e 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -9,6 +9,7 @@ add_subdirectory(minicpm5) add_subdirectory(qwen3) add_subdirectory(qwen3_5) add_subdirectory(lfm2) +add_subdirectory(ling3) add_subdirectory(qwen3_service) add_subdirectory(qwen3_moe) add_subdirectory(deepseek_ocr) diff --git a/examples/ling3/CMakeLists.txt b/examples/ling3/CMakeLists.txt new file mode 100644 index 000000000..a27d757ac --- /dev/null +++ b/examples/ling3/CMakeLists.txt @@ -0,0 +1,3 @@ +add_executable(mllm-ling3-runner main.cpp) +target_link_libraries(mllm-ling3-runner PRIVATE MllmRT MllmCPUBackend) +target_include_directories(mllm-ling3-runner PRIVATE ${MLLM_INCLUDE_DIR}) diff --git a/examples/ling3/README.md b/examples/ling3/README.md new file mode 100644 index 000000000..cbe6259e6 --- /dev/null +++ b/examples/ling3/README.md @@ -0,0 +1,70 @@ +# Ling-3.0-tiny mobile CPU runner + +This runner targets the pinned official checkpoint +`inclusionAI/Ling-3.0-tiny@a2ee06c0f2de5b171701aee7f73f70a1da75483b`. +The supported deployment envelope is batch 1, a 2048-token cache, float32 +recurrent state, and KAI W4A32 linear weights on ARM64 macOS and Android. +Build the ARM CPU backend without `-ffast-math`: Ling-3's recurrent gates are +not compatible with relaxed IEEE semantics. For Android, the validated ISA +flags are `-march=armv8.2-a+fp16+fp16fml+dotprod+i8mm`; do not append +`-ffast-math`. + +Validate the source checkpoint before conversion: + +```bash +python3 validate_checkpoint.py /path/to/Ling-3.0-tiny \ + --observed-revision a2ee06c0f2de5b171701aee7f73f70a1da75483b +``` + +Convert with the repository V2 converter and +`quant_cfg_tiny_w4a32_kai.json`, using model name `Ling-3.0-tiny`, then seal +the output descriptor table: + +```bash +python3 validate_converted_model.py /path/to/Ling-3.0-tiny.mllm \ + /path/to/Ling-3.0-tiny +``` + +Run one deterministic smoke request: + +```bash +./mllm-ling3-runner \ + --model_path /path/to/Ling-3.0-tiny.mllm \ + --tokenizer_path /path/to/Ling-3.0-tiny/tokenizer.json \ + --config_path config_tiny_w4a32_kai.json \ + --prompt '你好，请用一句话介绍你自己。' \ + --max_new_tokens 8 --print_token_ids +``` + +The runner emits `LING3_RUN_START`, generated token IDs, and +`LING3_RUN_OK`. A successful build or tokenizer-only test is not a full-model +runtime result; device evidence must retain the converted model SHA256 and +the runner/library identities together. + +For a longer correctness demo, pass the prompt directly and use a 64-token +generation limit: + +```bash +./mllm-ling3-runner \ + --model_path /path/to/Ling-3.0-tiny.mllm \ + --tokenizer_path /path/to/Ling-3.0-tiny/tokenizer.json \ + --config_path config_tiny_w4a32_kai.json \ + --prompt '请用中文详细介绍 Ling-3.0-tiny 的混合注意力架构，并解释 KDA、MLA 和 MoE 各自的作用。' --disable_thinking \ + --max_new_tokens 64 --print_token_ids +``` + +The expected completion marker is +`LING3_RUN_OK prompt_tokens=49 generated_tokens=64`. This is a generation +correctness demo, not a perplexity, model-quality, or performance benchmark. + +## Retained tests + +The Ling-specific coverage follows the repository test layout (`ctest` labels +in parentheses): + +| Layer | Target | Notes | +| --- | --- | --- | +| KDA kernel | `Mllm-Test-CPUKernel --gtest_filter='KimiDeltaAttentionKernelTest.*'` (`cpu-kernel`) | scalar-reference oracle, bitwise prefill/decode and serial/parallel checks | +| Public `nn::KimiDeltaAttention` | `Mllm-Test-Nn-KimiDeltaAttention` (`KimiDeltaAttentionFocused`) | eager reference, in-place state, chunked equivalence, trace + serialization | +| Causal convolution | upstream `Mllm-Test-Nn-CausalDepthwiseConv1D` and the `CausalDepthwiseConv*KernelTest` suites | Ling registers the upstream operation with the current-first order | +| Model | `Mllm-Test-Ling3-Config`, `Mllm-Test-Ling3-RoPE`, `Mllm-Test-Ling3-Tokenizer` (`ling3`) | set `MLLM_LING3_EXAMPLE_DIR` on a device; the tokenizer test skips unless `LING3_OFFICIAL_TOKENIZER` points at the official `tokenizer.json` | diff --git a/examples/ling3/config_tiny_w4a32_kai.json b/examples/ling3/config_tiny_w4a32_kai.json new file mode 100644 index 000000000..09e733d59 --- /dev/null +++ b/examples/ling3/config_tiny_w4a32_kai.json @@ -0,0 +1,46 @@ +{ + "architectures": ["BailingMoeV3ForCausalLM"], + "model_type": "bailing_hybrid", + "hidden_size": 1536, + "intermediate_size": 4608, + "num_hidden_layers": 24, + "num_attention_heads": 16, + "num_key_value_heads": 16, + "head_dim": 128, + "vocab_size": 157184, + "max_position_embeddings": 131072, + "rms_norm_eps": 1e-06, + "rope_theta": 6000000, + "hidden_act": "silu", + "tie_word_embeddings": false, + "use_qkv_bias": false, + "layer_group_size": 4, + "short_conv_kernel_size": 4, + "no_kda_lora": true, + "kda_safe_gate": true, + "kda_lower_bound": -5, + "q_lora_rank": 256, + "kv_lora_rank": 512, + "qk_rope_head_dim": 64, + "qk_nope_head_dim": 128, + "qk_head_dim": 192, + "v_head_dim": 128, + "rope_interleave": true, + "gated_attention_proj_granularity_type": "head_wise", + "num_experts": 128, + "num_shared_experts": 1, + "num_experts_per_tok": 8, + "n_group": 8, + "topk_group": 4, + "moe_intermediate_size": 512, + "moe_shared_expert_intermediate_size": 512, + "first_k_dense_replace": 1, + "routed_scaling_factor": 2.5, + "scoring_func": "sigmoid", + "topk_method": "noaux_tc", + "moe_router_enable_expert_bias": true, + "pad_token_id": 156892, + "eos_token_id": 156895, + "max_cache_length": 2048, + "linear_impl_type": "KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32" +} diff --git a/examples/ling3/main.cpp b/examples/ling3/main.cpp new file mode 100644 index 000000000..b9e7009e9 --- /dev/null +++ b/examples/ling3/main.cpp @@ -0,0 +1,124 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +using mllm::Argparse; + +namespace { + +std::string readPromptFile(const std::string& path) { + std::ifstream stream(path, std::ios::binary); + if (!stream) { throw std::invalid_argument("unable to read prompt_file: " + path); } + std::string text{std::istreambuf_iterator(stream), std::istreambuf_iterator()}; + while (!text.empty() && (text.back() == '\n' || text.back() == '\r')) { text.pop_back(); } + if (text.empty()) { throw std::invalid_argument("prompt_file must not be empty: " + path); } + return text; +} + +} // 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("MLLM V2 model path").required(true); + auto& tokenizer_path = Argparse::add("-t|--tokenizer_path").help("Official tokenizer.json").required(true); + auto& config_path = Argparse::add("-c|--config_path").help("Ling-3 mobile runtime config").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 a prompt from a UTF-8 file").required(false); + auto& max_new_tokens = + Argparse::add("-g|--max_new_tokens").help("Maximum generated tokens (default: 8)").required(false); + auto& min_new_tokens = Argparse::add("--min_new_tokens").help("Suppress EOS until this many tokens").required(false); + auto& disable_thinking = + Argparse::add("--disable_thinking").help("Use the official thinking-off chat template").required(false); + auto& print_token_ids = Argparse::add("--print_token_ids").help("Print generated token IDs to stderr").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; + + const auto config = mllm::models::ling3::Ling3Config(config_path.get()); + int generation_limit = max_new_tokens.isSet() ? max_new_tokens.get() : 8; + int minimum_generation = min_new_tokens.isSet() ? min_new_tokens.get() : 0; + if (generation_limit <= 0 || generation_limit > config.max_cache_length || minimum_generation < 0 + || minimum_generation > generation_limit) { + throw std::invalid_argument("generation lengths must satisfy 0 <= min_new_tokens <= max_new_tokens <= max_cache_length"); + } + if (prompt.isSet() && prompt_file.isSet()) { throw std::invalid_argument("prompt and prompt_file are mutually exclusive"); } + + std::string configured_prompt; + if (prompt_file.isSet()) { + configured_prompt = readPromptFile(prompt_file.get()); + } else if (prompt.isSet()) { + configured_prompt = prompt.get(); + } + + const auto parameters = mllm::load(model_path.get(), mllm::ModelFileVersion::kV2); + mllm::models::ling3::validateLing3ModelConfigMatch(config, parameters); + auto tokenizer = mllm::models::ling3::Ling3Tokenizer(tokenizer_path.get()); + auto model = mllm::models::ling3::Ling3ForCausalLM(config); + model.load(parameters); + fmt::print("Ling-3.0-tiny: {} layers ({} MLA + {} KDA), CPU threads={}\n", config.num_hidden_layers, + config.numFullAttentionLayers(), config.numKDALayers(), mllm::Context::instance().getCpuOpThreads()); + + int exit_code = 0; + while (true) { + std::string prompt_text = configured_prompt; + if (!prompt.isSet() && !prompt_file.isSet()) { + fmt::print("Prompt text (or 'exit/quit'): "); + if (!std::getline(std::cin, prompt_text) || prompt_text == "exit" || prompt_text == "quit") { break; } + } + if (prompt_text.empty()) { + if (prompt.isSet() || prompt_file.isSet()) { throw std::invalid_argument("prompt must not be empty"); } + continue; + } + try { + auto input = tokenizer.convertMessage({.prompt = prompt_text, + .system_prompt = "", + .enable_thinking = !(disable_thinking.isSet() && disable_thinking.get())}); + const int prompt_tokens = input.at("sequence").shape()[1]; + if (prompt_tokens + generation_limit - 1 > config.max_cache_length) { + throw std::invalid_argument("prompt plus generation exceeds max_cache_length"); + } + model.resetState(); + fmt::print("LING3_RUN_START prompt_tokens={} max_new_tokens={} min_new_tokens={}\nResponse: ", prompt_tokens, + generation_limit, minimum_generation); + int generated_tokens = 0; + for (const auto& step : model.chat(input, {{"max_length", mllm::AnyValue(generation_limit)}, + {"min_new_tokens", mllm::AnyValue(minimum_generation)}, + {"do_sample", mllm::AnyValue(false)}})) { + if (print_token_ids.isSet() && print_token_ids.get()) { fmt::print(stderr, "LING3_TOKEN_ID:{}\n", step.cur_token_id); } + fmt::print("{}", tokenizer.detokenizeBytes(step.cur_token_id)); + std::fflush(stdout); + ++generated_tokens; + } + fmt::print("\nLING3_RUN_OK prompt_tokens={} generated_tokens={}\n", prompt_tokens, generated_tokens); + } catch (const std::exception& error) { + fmt::print(stderr, "LING3_RUN_ERROR:{}\n", error.what()); + exit_code = 1; + } + if (prompt.isSet() || prompt_file.isSet()) { break; } + } + + model.perfSummary(); + mllm::memoryReport(); + return exit_code; +}) diff --git a/examples/ling3/quant_cfg_tiny_w4a32_kai.json b/examples/ling3/quant_cfg_tiny_w4a32_kai.json new file mode 100644 index 000000000..66855abac --- /dev/null +++ b/examples/ling3/quant_cfg_tiny_w4a32_kai.json @@ -0,0 +1,47 @@ +{ + "^model\\.layers\\.\\d+\\.attention\\.(q_proj|k_proj|v_proj|f_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, 1536], "replace": true} + }, + "^model\\.layers\\.(?:[0-2]|[4-6]|8|9|10|12|13|14|16|17|18|20|21|22)\\.attention\\.g_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, 1536], "replace": true} + }, + "^model\\.layers\\.(?:3|7|11|15|19|23)\\.attention\\.g_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": [16, 1536], "replace": true} + }, + "^model\\.layers\\.\\d+\\.attention\\.b_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": [16, 1536], "replace": true} + }, + "^model\\.layers\\.\\d+\\.attention\\.o_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": [1536, 2048], "replace": true} + }, + "^model\\.layers\\.\\d+\\.attention\\.q_a_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": [256, 1536], "replace": true} + }, + "^model\\.layers\\.\\d+\\.attention\\.q_b_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": [3072, 256], "replace": true} + }, + "^model\\.layers\\.\\d+\\.attention\\.kv_a_proj_with_mqa\\.weight$": { + "hints": {"quant_method": "kai", "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", "kai_matmul_layout": "mxk_nxk", "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", "shape": [576, 1536], "replace": true} + }, + "^model\\.layers\\.\\d+\\.attention\\.kv_b_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": [4096, 512], "replace": true} + }, + "^model\\.layers\\.\\d+\\.attention\\.dense\\.weight$": { + "hints": {"quant_method": "kai", "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", "kai_matmul_layout": "mxk_nxk", "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", "shape": [1536, 2048], "replace": true} + }, + "^model\\.layers\\.0\\.mlp\\.(gate_proj|up_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": [4608, 1536], "replace": true} + }, + "^model\\.layers\\.0\\.mlp\\.down_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": [1536, 4608], "replace": true} + }, + "^model\\.layers\\.(?:[1-9]|1\\d|2[0-3])\\.mlp\\.(?:experts\\.\\d+|shared_experts)\\.(gate_proj|up_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, 1536], "replace": true} + }, + "^model\\.layers\\.(?:[1-9]|1\\d|2[0-3])\\.mlp\\.(?:experts\\.\\d+|shared_experts)\\.down_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": [1536, 512], "replace": true} + }, + "^lm_head\\.weight$": { + "hints": {"quant_method": "kai", "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", "kai_matmul_layout": "mxk_nxk", "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", "shape": [157184, 1536], "replace": true} + } +} diff --git a/examples/ling3/test_validators.py b/examples/ling3/test_validators.py new file mode 100644 index 000000000..331a426de --- /dev/null +++ b/examples/ling3/test_validators.py @@ -0,0 +1,39 @@ +# Copyright (c) MLLM Team. +# Licensed under the MIT License. + +import json +import unittest +from pathlib import Path + +from validate_checkpoint import ( + FULL_ATTENTION_LAYERS, + expected_shapes, + validate_config, + validate_recipe, +) + + +HERE = Path(__file__).resolve().parent + + +class Ling3ValidatorTests(unittest.TestCase): + def test_expected_checkpoint_inventory_is_exact(self): + shapes = expected_shapes() + self.assertEqual(len(shapes), 9283) + self.assertEqual( + {layer for layer in range(24) if f"model.layers.{layer}.attention.q_a_proj.weight" in shapes}, + FULL_ATTENTION_LAYERS, + ) + + def test_runtime_and_recipe_match_official_contract(self): + runtime = json.loads((HERE / "config_tiny_w4a32_kai.json").read_text()) + validate_config(runtime, runtime) + recipe = json.loads((HERE / "quant_cfg_tiny_w4a32_kai.json").read_text()) + matched = validate_recipe(expected_shapes(), recipe) + self.assertGreater(len(matched), 9000) + self.assertNotIn("model.word_embeddings.weight", matched) + self.assertNotIn("model.layers.1.mlp.gate.weight", matched) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/ling3/validate_checkpoint.py b/examples/ling3/validate_checkpoint.py new file mode 100644 index 000000000..89ca6f560 --- /dev/null +++ b/examples/ling3/validate_checkpoint.py @@ -0,0 +1,274 @@ +#!/usr/bin/env python3 +# Copyright (c) MLLM Team. +# Licensed under the MIT License. + +"""Validate the pinned Ling-3.0-tiny source contract and W4A32 recipe.""" + +from __future__ import annotations + +import argparse +import json +import re +import struct +from pathlib import Path + + +PINNED_MODEL_ID = "inclusionAI/Ling-3.0-tiny" +PINNED_REVISION = "a2ee06c0f2de5b171701aee7f73f70a1da75483b" +FULL_ATTENTION_LAYERS = {3, 7, 11, 15, 19, 23} +OFFICIAL_CONTRACT = { + "architectures": ["BailingMoeV3ForCausalLM"], + "model_type": "bailing_hybrid", + "hidden_size": 1536, + "intermediate_size": 4608, + "num_hidden_layers": 24, + "num_attention_heads": 16, + "num_key_value_heads": 16, + "head_dim": 128, + "vocab_size": 157184, + "max_position_embeddings": 131072, + "rms_norm_eps": 1e-6, + "rope_theta": 6000000, + "layer_group_size": 4, + "short_conv_kernel_size": 4, + "no_kda_lora": True, + "kda_safe_gate": True, + "kda_lower_bound": -5, + "q_lora_rank": 256, + "kv_lora_rank": 512, + "qk_rope_head_dim": 64, + "qk_nope_head_dim": 128, + "qk_head_dim": 192, + "v_head_dim": 128, + "rope_interleave": True, + "gated_attention_proj_granularity_type": "head_wise", + "num_experts": 128, + "num_shared_experts": 1, + "num_experts_per_tok": 8, + "n_group": 8, + "topk_group": 4, + "moe_intermediate_size": 512, + "moe_shared_expert_intermediate_size": 512, + "first_k_dense_replace": 1, + "routed_scaling_factor": 2.5, + "scoring_func": "sigmoid", + "topk_method": "noaux_tc", + "moe_router_enable_expert_bias": True, + "tie_word_embeddings": False, + "use_qkv_bias": False, + "pad_token_id": 156892, + "eos_token_id": 156895, +} +KAI_HINTS = { + "quant_method": "kai", + "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", + "kai_matmul_layout": "mxk_nxk", + "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", +} + + +def expected_shapes() -> dict[str, list[int]]: + shapes: dict[str, list[int]] = { + "model.word_embeddings.weight": [157184, 1536], + "model.norm.weight": [1536], + "lm_head.weight": [157184, 1536], + } + for layer in range(24): + prefix = f"model.layers.{layer}" + shapes[f"{prefix}.input_layernorm.weight"] = [1536] + shapes[f"{prefix}.post_attention_layernorm.weight"] = [1536] + attention = f"{prefix}.attention" + if layer in FULL_ATTENTION_LAYERS: + shapes.update( + { + f"{attention}.q_a_proj.weight": [256, 1536], + f"{attention}.q_a_layernorm.weight": [256], + f"{attention}.q_b_proj.weight": [3072, 256], + f"{attention}.kv_a_proj_with_mqa.weight": [576, 1536], + f"{attention}.kv_a_layernorm.weight": [512], + f"{attention}.kv_b_proj.weight": [4096, 512], + f"{attention}.g_proj.weight": [16, 1536], + f"{attention}.dense.weight": [1536, 2048], + } + ) + else: + shapes.update( + { + f"{attention}.q_proj.weight": [2048, 1536], + f"{attention}.k_proj.weight": [2048, 1536], + f"{attention}.v_proj.weight": [2048, 1536], + f"{attention}.q_conv1d.weight": [2048, 1, 4], + f"{attention}.k_conv1d.weight": [2048, 1, 4], + f"{attention}.v_conv1d.weight": [2048, 1, 4], + f"{attention}.f_proj.weight": [2048, 1536], + f"{attention}.A_log": [16], + f"{attention}.dt_bias": [2048], + f"{attention}.b_proj.weight": [16, 1536], + f"{attention}.g_proj.weight": [2048, 1536], + f"{attention}.o_norm.weight": [128], + f"{attention}.o_proj.weight": [1536, 2048], + } + ) + mlp = f"{prefix}.mlp" + if layer == 0: + shapes[f"{mlp}.gate_proj.weight"] = [4608, 1536] + shapes[f"{mlp}.up_proj.weight"] = [4608, 1536] + shapes[f"{mlp}.down_proj.weight"] = [1536, 4608] + else: + shapes[f"{mlp}.gate.weight"] = [128, 1536] + shapes[f"{mlp}.gate.expert_bias"] = [128] + for expert in range(128): + expert_prefix = f"{mlp}.experts.{expert}" + shapes[f"{expert_prefix}.gate_proj.weight"] = [512, 1536] + shapes[f"{expert_prefix}.up_proj.weight"] = [512, 1536] + shapes[f"{expert_prefix}.down_proj.weight"] = [1536, 512] + shapes[f"{mlp}.shared_experts.gate_proj.weight"] = [512, 1536] + shapes[f"{mlp}.shared_experts.up_proj.weight"] = [512, 1536] + shapes[f"{mlp}.shared_experts.down_proj.weight"] = [1536, 512] + assert len(shapes) == 9283 + return shapes + + +def validate_config(checkpoint_config: dict, runtime_config: dict) -> None: + mismatches = [ + f"{key}={checkpoint_config.get(key)!r}, expected {value!r}" + for key, value in OFFICIAL_CONTRACT.items() + if checkpoint_config.get(key) != value + ] + if mismatches: + raise AssertionError("Checkpoint contract mismatch: " + "; ".join(mismatches)) + runtime_mismatches = [ + f"{key}={runtime_config.get(key)!r}, expected {value!r}" + for key, value in OFFICIAL_CONTRACT.items() + if runtime_config.get(key) != value + ] + if runtime_mismatches: + raise AssertionError("Runtime contract mismatch: " + "; ".join(runtime_mismatches)) + if runtime_config.get("max_cache_length") != 2048: + raise AssertionError("Mobile runtime must use the reviewed 2048-token cache bound") + expected_impl = ( + "KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_" + "qai8dxp1x8_qsi4c32p8x8_1x8x32" + ) + if runtime_config.get("linear_impl_type") != expected_impl: + raise AssertionError("Runtime config does not select the reviewed KAI W4A32 path") + + +def validate_recipe(shapes: dict[str, list[int]], recipe: dict) -> set[str]: + matched: set[str] = set() + for pattern_text, entry in recipe.items(): + pattern = re.compile(pattern_text) + names = {name for name in shapes if pattern.fullmatch(name)} + if not names: + raise AssertionError(f"Quantization pattern matched no tensors: {pattern_text}") + overlap = matched & names + if overlap: + raise AssertionError(f"Quantization patterns overlap: {sorted(overlap)[:5]}") + hints = entry.get("hints", {}) + for key, value in KAI_HINTS.items(): + if hints.get(key) != value: + raise AssertionError(f"{pattern_text}: invalid {key}") + if hints.get("replace") is not True: + raise AssertionError(f"{pattern_text}: KAI weights must replace source weights") + for name in names: + if shapes[name] != hints.get("shape"): + raise AssertionError( + f"{pattern_text}: {name} has shape {shapes[name]}, recipe says {hints.get('shape')}" + ) + matched.update(names) + + expected_quantized = { + name + for name, shape in shapes.items() + if name == "lm_head.weight" + or ( + name.endswith(".weight") + and len(shape) == 2 + and name != "model.word_embeddings.weight" + and not name.endswith(".mlp.gate.weight") + ) + } + if matched != expected_quantized: + raise AssertionError( + "Quantized tensor coverage mismatch: " + f"missing={sorted(expected_quantized - matched)[:8]}, " + f"extra={sorted(matched - expected_quantized)[:8]}" + ) + return matched + + +def read_safetensors_header(path: Path) -> dict: + with path.open("rb") as file: + raw_length = file.read(8) + if len(raw_length) != 8: + raise AssertionError(f"Truncated safetensors header: {path}") + header_length = struct.unpack(" None: + actual: dict[str, tuple[str, list[int]]] = {} + for shard_name in sorted(set(weight_map.values())): + header = read_safetensors_header(checkpoint / shard_name) + for name, descriptor in header.items(): + if name == "__metadata__": + continue + if name in actual: + raise AssertionError(f"Duplicate tensor across shards: {name}") + actual[name] = (descriptor["dtype"], descriptor["shape"]) + if set(actual) != set(shapes): + raise AssertionError( + f"Shard tensor set mismatch: missing={sorted(set(shapes) - set(actual))[:8]}, " + f"extra={sorted(set(actual) - set(shapes))[:8]}" + ) + for name, expected_shape in shapes.items(): + expected_dtype = ( + "F32" + if name.endswith(".A_log") + or name.endswith(".dt_bias") + or name.endswith(".gate.expert_bias") + else "BF16" + ) + if actual[name] != (expected_dtype, expected_shape): + raise AssertionError(f"{name}: expected {(expected_dtype, expected_shape)}, got {actual[name]}") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint", type=Path) + parser.add_argument("--runtime-config", type=Path, default=Path(__file__).with_name("config_tiny_w4a32_kai.json")) + parser.add_argument("--quant-config", type=Path, default=Path(__file__).with_name("quant_cfg_tiny_w4a32_kai.json")) + parser.add_argument("--observed-revision", required=True) + parser.add_argument("--index-only", action="store_true") + args = parser.parse_args() + if args.observed_revision != PINNED_REVISION: + raise AssertionError( + f"Checkpoint revision {args.observed_revision} does not match pinned {PINNED_REVISION}" + ) + 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()) + validate_config(checkpoint_config, runtime_config) + shapes = expected_shapes() + index = json.loads((args.checkpoint / "model.safetensors.index.json").read_text()) + weight_map = index.get("weight_map", {}) + if set(weight_map) != set(shapes): + raise AssertionError( + f"Index tensor set mismatch: missing={sorted(set(shapes) - set(weight_map))[:8]}, " + f"extra={sorted(set(weight_map) - set(shapes))[:8]}" + ) + validate_recipe(shapes, recipe) + if not args.index_only: + validate_shards(args.checkpoint, weight_map, shapes) + print( + "LING3_CHECKPOINT_OK " + f"model={PINNED_MODEL_ID} revision={PINNED_REVISION} tensors={len(shapes)} " + f"shards={len(set(weight_map.values()))} mode={'index' if args.index_only else 'full'}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/ling3/validate_converted_model.py b/examples/ling3/validate_converted_model.py new file mode 100644 index 000000000..ce1ea92d4 --- /dev/null +++ b/examples/ling3/validate_converted_model.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +# Copyright (c) MLLM Team. +# Licensed under the MIT License. + +"""Validate a converted Ling-3.0-tiny MLLM V2 model 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 + + +MODEL_HEADER = struct.Struct(" int: + if out_channels <= 0 or in_channels <= 0 or in_channels % 32: + raise AssertionError("KAI W4A32 shapes require positive dimensions and K divisible by 32") + nr = 8 + blocks_per_row = in_channels // 32 + bytes_per_nr_rows = nr * (blocks_per_row * (16 + 2) + 4 + 4) + return math.ceil(out_channels / nr) * bytes_per_nr_rows + + +def decode_c_string(value: bytes) -> str: + return value.split(b"\0", 1)[0].decode("utf-8") + + +def expected_descriptors(shapes: dict[str, list[int]], recipe: dict) -> dict[str, tuple[int, list[int], int]]: + expected = {name: (FLOAT32, shape, math.prod(shape) * 4) for name, shape in shapes.items()} + validate_recipe(shapes, recipe) + for pattern_text, entry in recipe.items(): + pattern = re.compile(pattern_text) + for name, shape in shapes.items(): + if pattern.fullmatch(name): + size = packed_size(shape[0], shape[1]) + expected[name] = (BYTE, [size], size) + return expected + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("model", type=Path) + parser.add_argument("checkpoint", type=Path) + parser.add_argument("--runtime-config", type=Path, default=Path(__file__).with_name("config_tiny_w4a32_kai.json")) + parser.add_argument("--quant-config", type=Path, default=Path(__file__).with_name("quant_cfg_tiny_w4a32_kai.json")) + parser.add_argument("--model-name", default="Ling-3.0-tiny") + args = parser.parse_args() + source_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()) + validate_config(source_config, runtime_config) + expected = expected_descriptors(expected_shapes(), recipe) + + file_size = args.model.stat().st_size + with args.model.open("rb") as model_file: + raw_header = model_file.read(MODEL_HEADER.size) + if len(raw_header) != MODEL_HEADER.size: + raise AssertionError("Truncated MLLM model header") + magic, version, raw_name, num_parameters, descriptor_offset = MODEL_HEADER.unpack(raw_header) + identity = (magic, version, decode_c_string(raw_name), descriptor_offset) + expected_identity = (MODEL_MAGIC, MODEL_VERSION, args.model_name, MODEL_HEADER.size) + if identity != expected_identity: + raise AssertionError(f"Invalid model identity: expected {expected_identity}, got {identity}") + if num_parameters != len(expected): + raise AssertionError(f"Expected {len(expected)} parameters, model declares {num_parameters}") + actual: dict[str, tuple[int, list[int], int, int]] = {} + for expected_id in range(num_parameters): + raw_descriptor = model_file.read(PARAMETER_DESCRIPTOR.size) + if len(raw_descriptor) != PARAMETER_DESCRIPTOR.size: + raise AssertionError(f"Truncated descriptor {expected_id}") + values = PARAMETER_DESCRIPTOR.unpack(raw_descriptor) + parameter_id, dtype, size, offset, shape_length = values[:5] + if parameter_id != expected_id or shape_length > 16: + raise AssertionError(f"Invalid parameter descriptor {expected_id}") + name = decode_c_string(values[21]) + if name in actual: + raise AssertionError(f"Duplicate converted parameter: {name}") + actual[name] = (dtype, list(values[5:21])[:shape_length], size, offset) + + if set(actual) != set(expected): + raise AssertionError( + f"Converted tensor set mismatch: missing={sorted(set(expected) - set(actual))[:8]}, " + f"extra={sorted(set(actual) - set(expected))[:8]}" + ) + data_start = MODEL_HEADER.size + len(expected) * PARAMETER_DESCRIPTOR.size + next_offset = data_start + for name, descriptor in sorted(actual.items(), key=lambda item: item[1][3]): + dtype, shape, size, offset = descriptor + if (dtype, shape, size) != expected[name]: + raise AssertionError(f"{name}: expected {expected[name]}, got {(dtype, shape, size)}") + if offset != next_offset: + raise AssertionError(f"{name}: expected offset {next_offset}, got {offset}") + next_offset += size + if next_offset != file_size: + raise AssertionError(f"Converted data ends at {next_offset}, file size is {file_size}") + quantized = sum(dtype == BYTE for dtype, _, _, _ in actual.values()) + print( + f"LING3_CONVERTED_MODEL_OK model={args.model} tensors={len(actual)} " + f"kai_w4a32={quantized} bytes={file_size}" + ) + + +if __name__ == "__main__": + main() diff --git a/mllm/backends/cpu/CPUBackend.cpp b/mllm/backends/cpu/CPUBackend.cpp index 190f61109..c984eea8c 100644 --- a/mllm/backends/cpu/CPUBackend.cpp +++ b/mllm/backends/cpu/CPUBackend.cpp @@ -40,6 +40,7 @@ #include "mllm/backends/cpu/ops/ISTFTOp.hpp" #include "mllm/backends/cpu/ops/IndexOp.hpp" #include "mllm/backends/cpu/ops/KVCacheOp.hpp" +#include "mllm/backends/cpu/ops/KimiDeltaAttentionOp.hpp" #include "mllm/backends/cpu/ops/LayerNormOp.hpp" #include "mllm/backends/cpu/ops/LinearOp.hpp" #include "mllm/backends/cpu/ops/MatMulOp.hpp" @@ -88,7 +89,7 @@ CPUBackend::CPUBackend() : Backend(kCPU, createCPUAllocator()) { CPUArgsortOpFactory, CPUCloneOpFactory, CPUAvgPool1dOpFactory, CPUFlashAttention2SwaSinkOpFactory, CPURadixAttnRelaxOpFactory, CPURadixAttnSwaSinkOpFactory, CPUEqualOpFactory, CPUWhereOpFactory, CPUGatherOpFactory, CPUCausalDepthwiseConv1DOpFactory, CPUGatedDeltaRuleOpFactory, - CPUGroupedQueryAttentionOpFactory, CPUParallelLinearOpFactory>(); + CPUGroupedQueryAttentionOpFactory, CPUParallelLinearOpFactory, CPUKimiDeltaAttentionOpFactory>(); } CPUBackend::~CPUBackend() { diff --git a/mllm/backends/cpu/kernels/common/kda/kimi_delta_attention.cpp b/mllm/backends/cpu/kernels/common/kda/kimi_delta_attention.cpp new file mode 100644 index 000000000..1609a2281 --- /dev/null +++ b/mllm/backends/cpu/kernels/common/kda/kimi_delta_attention.cpp @@ -0,0 +1,195 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/kernels/common/kda/kimi_delta_attention.hpp" + +#include "mllm/engine/Context.hpp" +#include "mllm/core/Parallel.hpp" + +#include +#include +#include +#include +#include + +#if defined(__linux__) || defined(__ANDROID__) +#include +#endif + +#if defined(__aarch64__) +#include +#endif + +namespace mllm::cpu::kda { +namespace { + +constexpr int kMaxParallelLanes = 8; +constexpr std::size_t kMinParallelWork = 65536; + +int availableCpuCount(int fallback) { +#if defined(__linux__) || defined(__ANDROID__) + cpu_set_t affinity = {}; + if (sched_getaffinity(0, sizeof(affinity), &affinity) == 0) { return CPU_COUNT(&affinity); } +#endif + return fallback; +} + +float stableSigmoid(float value) { + if (value >= 0.0F) { + const float exp_value = std::exp(-value); + return 1.0F / (1.0F + exp_value); + } + const float exp_value = std::exp(value); + return exp_value / (1.0F + exp_value); +} + +float stableSoftplus(float value) { + if (value > 20.0F) { return value; } + if (value < -20.0F) { return std::exp(value); } + return std::log1p(std::exp(value)); +} + +float squaredNorm(const float* values, int length) { + float sum = 0.0F; + int index = 0; +#if defined(__aarch64__) + float32x4_t vector_sum = vdupq_n_f32(0.0F); + for (; index + 4 <= length; index += 4) { + const float32x4_t value = vld1q_f32(values + index); + vector_sum = vmlaq_f32(vector_sum, value, value); + } + sum = vaddvq_f32(vector_sum); +#endif + for (; index < length; ++index) { sum += values[index] * values[index]; } + return sum; +} + +void normalize(const float* source, float scale, float* destination, int length) { + int index = 0; +#if defined(__aarch64__) + for (; index + 4 <= length; index += 4) { vst1q_f32(destination + index, vmulq_n_f32(vld1q_f32(source + index), scale)); } +#endif + for (; index < length; ++index) { destination[index] = source[index] * scale; } +} + +void decayAndAccumulatePrediction(float* state_row, float decay, float key, float* prediction, int length) { + int index = 0; +#if defined(__aarch64__) + for (; index + 4 <= length; index += 4) { + float32x4_t state_value = vmulq_n_f32(vld1q_f32(state_row + index), decay); + vst1q_f32(state_row + index, state_value); + float32x4_t predicted = vld1q_f32(prediction + index); + predicted = vmlaq_n_f32(predicted, state_value, key); + vst1q_f32(prediction + index, predicted); + } +#endif + for (; index < length; ++index) { + state_row[index] *= decay; + prediction[index] += key * state_row[index]; + } +} + +void updateAndAccumulateOutput(float* state_row, float key, float query, const float* delta, float* output, int length) { + int index = 0; +#if defined(__aarch64__) + for (; index + 4 <= length; index += 4) { + float32x4_t state_value = vld1q_f32(state_row + index); + state_value = vmlaq_n_f32(state_value, vld1q_f32(delta + index), key); + vst1q_f32(state_row + index, state_value); + float32x4_t output_value = vld1q_f32(output + index); + output_value = vmlaq_n_f32(output_value, state_value, query); + vst1q_f32(output + index, output_value); + } +#endif + for (; index < length; ++index) { + state_row[index] += key * delta[index]; + output[index] += query * state_row[index]; + } +} + +} // namespace + +void kimiDeltaAttentionF32(const float* q, const float* k, const float* v, const float* gate_logits, const float* beta, + const float* a_log, const float* dt_bias, float* state, float* output, int batch_size, + int sequence_length, int num_heads, int head_dim, bool safe_gate, float lower_bound) { + kimiDeltaAttentionF32(q, k, v, gate_logits, beta, a_log, dt_bias, state, output, batch_size, sequence_length, num_heads, + head_dim, safe_gate, lower_bound, ::mllm::Context::instance().getCpuOpThreads()); +} + +void kimiDeltaAttentionF32(const float* q, const float* k, const float* v, const float* gate_logits, const float* beta, + const float* a_log, const float* dt_bias, float* state, float* output, int batch_size, + int sequence_length, int num_heads, int head_dim, bool safe_gate, float lower_bound, + int thread_count) { + if (q == nullptr || k == nullptr || v == nullptr || gate_logits == nullptr || beta == nullptr || a_log == nullptr + || dt_bias == nullptr || state == nullptr || output == nullptr) { + throw std::invalid_argument("Kimi Delta Attention received a null pointer"); + } + if (batch_size <= 0 || sequence_length <= 0 || num_heads <= 0 || head_dim <= 0 || thread_count <= 0) { + throw std::invalid_argument("Kimi Delta Attention received an invalid shape"); + } + if (safe_gate && (!std::isfinite(lower_bound) || lower_bound >= 0.0F)) { + throw std::invalid_argument("Kimi Delta Attention safe gate requires a finite negative lower bound"); + } + + const float query_dim_scale = 1.0F / std::sqrt(static_cast(head_dim)); + const int task_count = batch_size * num_heads; + const std::size_t work = static_cast(task_count) * sequence_length * head_dim * head_dim; + const int parallel_lanes = std::min({thread_count, task_count, availableCpuCount(thread_count), kMaxParallelLanes}); + const bool use_parallel = parallel_lanes > 1 && work >= kMinParallelWork; + const int scheduled_lanes = use_parallel ? parallel_lanes : 1; + + const auto run_lane = [&](int lane) { + std::vector normalized_query(static_cast(head_dim)); + std::vector normalized_key(static_cast(head_dim)); + std::vector decay(static_cast(head_dim)); + std::vector prediction(static_cast(head_dim)); + std::vector delta(static_cast(head_dim)); + + for (int task = lane; task < task_count; task += scheduled_lanes) { + const int batch = task / num_heads; + const int head = task % num_heads; + const std::size_t state_base = (static_cast(batch) * num_heads + head) * head_dim * head_dim; + const float a_scale = std::exp(a_log[head]); + + for (int token = 0; token < sequence_length; ++token) { + const std::size_t vector_base = + ((static_cast(batch) * sequence_length + token) * num_heads + head) * head_dim; + const std::size_t beta_index = (static_cast(batch) * sequence_length + token) * num_heads + head; + + const float query_scale = query_dim_scale / std::sqrt(squaredNorm(q + vector_base, head_dim) + 1.0e-6F); + const float key_scale = 1.0F / std::sqrt(squaredNorm(k + vector_base, head_dim) + 1.0e-6F); + normalize(q + vector_base, query_scale, normalized_query.data(), head_dim); + normalize(k + vector_base, key_scale, normalized_key.data(), head_dim); + + for (int key_dim = 0; key_dim < head_dim; ++key_dim) { + const float gate_input = gate_logits[vector_base + key_dim] + dt_bias[head * head_dim + key_dim]; + const float log_decay = + safe_gate ? lower_bound * stableSigmoid(a_scale * gate_input) : -a_scale * stableSoftplus(gate_input); + decay[key_dim] = std::exp(log_decay); + } + + std::fill(prediction.begin(), prediction.end(), 0.0F); + for (int key_dim = 0; key_dim < head_dim; ++key_dim) { + float* state_row = state + state_base + static_cast(key_dim) * head_dim; + decayAndAccumulatePrediction(state_row, decay[key_dim], normalized_key[key_dim], prediction.data(), head_dim); + } + for (int value_dim = 0; value_dim < head_dim; ++value_dim) { + delta[value_dim] = beta[beta_index] * (v[vector_base + value_dim] - prediction[value_dim]); + } + + float* token_output = output + vector_base; + std::fill(token_output, token_output + head_dim, 0.0F); + for (int key_dim = 0; key_dim < head_dim; ++key_dim) { + float* state_row = state + state_base + static_cast(key_dim) * head_dim; + updateAndAccumulateOutput(state_row, normalized_key[key_dim], normalized_query[key_dim], delta.data(), token_output, + head_dim); + } + } + } + }; + + MLLM_CONDITIONAL_PARALLEL_FOR(use_parallel, scheduled_lanes, lane, 0, scheduled_lanes, 1, + { run_lane(static_cast(lane)); }); +} + +} // namespace mllm::cpu::kda diff --git a/mllm/backends/cpu/kernels/common/kda/kimi_delta_attention.hpp b/mllm/backends/cpu/kernels/common/kda/kimi_delta_attention.hpp new file mode 100644 index 000000000..38a9ac32a --- /dev/null +++ b/mllm/backends/cpu/kernels/common/kda/kimi_delta_attention.hpp @@ -0,0 +1,28 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +namespace mllm::cpu::kda { + +// Kimi Delta Attention recurrent update used by Ling-3.0. +// +// q/k/v/g: [B, S, H, D] +// beta: [B, S, H] (already passed through sigmoid) +// a_log: [H] +// dt_bias: [H, D] +// state: [B, H, D, D], updated in place +// output: [B, S, H, D] +// +// q and k are L2-normalized inside this function and q is additionally +// scaled by 1 / sqrt(D), matching the FLA KDA reference implementation. +void kimiDeltaAttentionF32(const float* q, const float* k, const float* v, const float* gate_logits, const float* beta, + const float* a_log, const float* dt_bias, float* state, float* output, int batch_size, + int sequence_length, int num_heads, int head_dim, bool safe_gate, float lower_bound); + +void kimiDeltaAttentionF32(const float* q, const float* k, const float* v, const float* gate_logits, const float* beta, + const float* a_log, const float* dt_bias, float* state, float* output, int batch_size, + int sequence_length, int num_heads, int head_dim, bool safe_gate, float lower_bound, + int thread_count); + +} // namespace mllm::cpu::kda diff --git a/mllm/backends/cpu/kernels/common/paged_attn/arch.hpp b/mllm/backends/cpu/kernels/common/paged_attn/arch.hpp index 4abb818d5..3e0901518 100644 --- a/mllm/backends/cpu/kernels/common/paged_attn/arch.hpp +++ b/mllm/backends/cpu/kernels/common/paged_attn/arch.hpp @@ -14,6 +14,7 @@ #include #include "mllm/backends/cpu/kernels/arm/math.hpp" #elif defined(MLLM_HOST_ARCH_X86_64) || defined(MLLM_HOST_ARCH_X86) +#include #else #endif @@ -36,9 +37,8 @@ struct VectorDotProduct { template<> struct VectorDotProduct { __MLLM_UNSAFE_OPT_BEGIN_O3 // Do not open fast math here - static MLLM_FORCE_INLINE void - run(const mllm_fp32_t* __restrict__ __lhs, const mllm_fp32_t* __restrict__ __rhs, mllm_fp32_t* __restrict__ __out, - size_t len) { + static MLLM_FORCE_INLINE void run(const mllm_fp32_t* __restrict__ __lhs, const mllm_fp32_t* __restrict__ __rhs, + mllm_fp32_t* __restrict__ __out, size_t len) { mllm_fp32_t sum = 0.0f; for (size_t i = 0; i < len; ++i) { sum += __lhs[i] * __rhs[i]; } *__out = sum; diff --git a/mllm/backends/cpu/ops/KimiDeltaAttentionOp.cpp b/mllm/backends/cpu/ops/KimiDeltaAttentionOp.cpp new file mode 100644 index 000000000..8d9b76177 --- /dev/null +++ b/mllm/backends/cpu/ops/KimiDeltaAttentionOp.cpp @@ -0,0 +1,31 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/ops/KimiDeltaAttentionOp.hpp" + +#include +#include + +#include "mllm/backends/cpu/kernels/common/kda/kimi_delta_attention.hpp" + +namespace mllm::cpu { + +CPUKimiDeltaAttentionOp::CPUKimiDeltaAttentionOp(const aops::KimiDeltaAttentionOpOptions& options) + : aops::KimiDeltaAttentionOp(options) {} + +void CPUKimiDeltaAttentionOp::forward(const std::vector& inputs, std::vector& outputs) { + for (const auto& input : inputs) { + if (!input.isContiguous()) { throw std::invalid_argument("KimiDeltaAttention CPU inputs must be contiguous"); } + } + + const auto& q = inputs[0]; + auto& output = outputs[0]; + auto& updated_state = outputs[1]; + if (!options_.state_inplace) { std::memcpy(updated_state.ptr(), inputs[7].ptr(), inputs[7].bytes()); } + kda::kimiDeltaAttentionF32(inputs[0].ptr(), inputs[1].ptr(), inputs[2].ptr(), inputs[3].ptr(), + inputs[4].ptr(), inputs[5].ptr(), inputs[6].ptr(), updated_state.ptr(), + output.ptr(), q.shape()[0], q.shape()[1], q.shape()[2], q.shape()[3], options_.safe_gate, + options_.lower_bound, options_.getThreads()); +} + +} // namespace mllm::cpu diff --git a/mllm/backends/cpu/ops/KimiDeltaAttentionOp.hpp b/mllm/backends/cpu/ops/KimiDeltaAttentionOp.hpp new file mode 100644 index 000000000..406c5af8f --- /dev/null +++ b/mllm/backends/cpu/ops/KimiDeltaAttentionOp.hpp @@ -0,0 +1,24 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include "mllm/core/aops/KimiDeltaAttentionOp.hpp" + +namespace mllm::cpu { + +class CPUKimiDeltaAttentionOp final : public aops::KimiDeltaAttentionOp { + public: + explicit CPUKimiDeltaAttentionOp(const aops::KimiDeltaAttentionOpOptions& options); + + void forward(const std::vector& inputs, std::vector& outputs) override; +}; + +class CPUKimiDeltaAttentionOpFactory : public TypedOpFactory { + public: + std::shared_ptr createOpImpl(const aops::KimiDeltaAttentionOpOptions& 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 7114ef549..891ec5441 100644 --- a/mllm/compile/ir/GeneratedRTTIKind.hpp +++ b/mllm/compile/ir/GeneratedRTTIKind.hpp @@ -1,4 +1,4 @@ -// Auto generated: 2026-08-21 15:20:12 +// Auto generated: 2026-09-06 09:09:14 // do not modify this file #pragma once @@ -12,159 +12,162 @@ namespace mllm::ir { // // RTTI Kind (RK) enum NodeKind : uint32_t { - RK_None, - RK_Op, - RK_Op_LinalgIROp, - RK_Op_LinalgIROp_RegisterOp, - RK_Op_LinalgIROp_CustomKernelOp, - RK_Op_LinalgIROp_FillOp, - RK_Op_LinalgIROp_AddOp, - RK_Op_LinalgIROp_SubOp, - RK_Op_LinalgIROp_STFTOp, - RK_Op_LinalgIROp_ISTFTOp, - RK_Op_LinalgIROp_MulOp, - RK_Op_LinalgIROp_DivOp, - RK_Op_LinalgIROp_AbsOp, - RK_Op_LinalgIROp_LogOp, - RK_Op_LinalgIROp_MatMulOp, - RK_Op_LinalgIROp_EmbeddingOp, - RK_Op_LinalgIROp_LinearOp, - RK_Op_LinalgIROp_RoPEOp, - RK_Op_LinalgIROp_SoftmaxOp, - RK_Op_LinalgIROp_TransposeOp, - RK_Op_LinalgIROp_RMSNormOp, - RK_Op_LinalgIROp_SiLUOp, - RK_Op_LinalgIROp_KVCacheOp, - RK_Op_LinalgIROp_CausalMaskOp, - RK_Op_LinalgIROp_CastTypeOp, - RK_Op_LinalgIROp_X2XOp, - RK_Op_LinalgIROp_ViewOp, - RK_Op_LinalgIROp_SplitOp, - RK_Op_LinalgIROp_FlashAttention2Op, - RK_Op_LinalgIROp_CausalDepthwiseConv1DOp, - RK_Op_LinalgIROp_GroupedQueryAttentionOp, - RK_Op_LinalgIROp_ParallelLinearOp, - RK_Op_LinalgIROp_GatedDeltaRuleOp, - RK_Op_LinalgIROp_RepeatOp, - RK_Op_LinalgIROp_PermuteOp, - RK_Op_LinalgIROp_Conv1DOp, - RK_Op_LinalgIROp_Conv2DOp, - RK_Op_LinalgIROp_Conv3DOp, - RK_Op_LinalgIROp_GELUOp, - RK_Op_LinalgIROp_LayerNormOp, - RK_Op_LinalgIROp_MultimodalRoPEOp, - RK_Op_LinalgIROp_VisionRoPEOp, - RK_Op_LinalgIROp_QuickGELUOp, - RK_Op_LinalgIROp_CopyOp, - RK_Op_LinalgIROp_CloneOp, - RK_Op_LinalgIROp_NegOp, - RK_Op_LinalgIROp_ConcatOp, - RK_Op_LinalgIROp_ReduceMinOp, - RK_Op_LinalgIROp_ReduceMaxOp, - RK_Op_LinalgIROp_ReduceSumOp, - RK_Op_LinalgIROp_ReLUOp, - RK_Op_LinalgIROp_ContiguousOp, - RK_Op_LinalgIROp_ReshapeOp, - RK_Op_LinalgIROp_SliceOp, - RK_Op_LinalgIROp_ParamOp, - RK_Op_LinalgIROp_IndexOp, - RK_Op_LinalgIROp_TopKOp, - RK_Op_LinalgIROp_MeanOp, - RK_Op_LinalgIROp_ClipOp, - RK_Op_LinalgIROp_ExpOp, - RK_Op_LinalgIROp_SinOp, - RK_Op_LinalgIROp_CosOp, - RK_Op_LinalgIROp_PagedAttnOp, - RK_Op_LinalgIROp_LayerNorm2DOp, - RK_Op_LinalgIROp_PadOp, - RK_Op_LinalgIROp_InterpolateOp, - RK_Op_LinalgIROp_EinsumOp, - RK_Op_LinalgIROp_StackOp, - RK_Op_LinalgIROp_MaskedScatterOp, - RK_Op_LinalgIROp_ScatterOp, - RK_Op_LinalgIROp_GatherOp, - RK_Op_LinalgIROp_ArgsortOp, - RK_Op_LinalgIROp_FlashAttention2SwaSinkOp, - RK_Op_LinalgIROp_RadixAttnRelaxOp, - RK_Op_LinalgIROp_RadixAttnSwaSinkOp, - RK_Op_LinalgIROp_EqualOp, - RK_Op_LinalgIROp_WhereOp, - RK_Op_LinalgIROp_SigmoidOp, - RK_Op_LinalgIROp_CustomizedOp, - RK_Op_LinalgIROp_Last, - RK_Op_GraphIROp, - RK_Op_GraphIROp_SubGraphOp, - RK_Op_GraphIROp_CallGraphOp, - RK_Op_GraphIROp_Last, - RK_Op_TensorIROp, - RK_Op_TensorIROp_AllocOp, - RK_Op_TensorIROp_RegisterOp, - RK_Op_TensorIROp_FreeOp, - RK_Op_TensorIROp_Last, - RK_Op_BuiltinIROp, - RK_Op_BuiltinIROp_ModuleOp, - RK_Op_BuiltinIROp_Last, - RK_Op_ControlFlowIROp, - RK_Op_ControlFlowIROp_ReturnOp, - RK_Op_ControlFlowIROp_Last, - RK_Op_ProgramIROp, - RK_Op_ProgramIROp_FragmentOp, - RK_Op_ProgramIROp_KernelLaunchOp, - RK_Op_ProgramIROp_KernelSymbolOp, - RK_Op_ProgramIROp_ValueSymbolOp, - RK_Op_ProgramIROp_JumpOp, - RK_Op_ProgramIROp_LabelOp, - RK_Op_ProgramIROp_ExitOp, - RK_Op_ProgramIROp_RetOp, - RK_Op_ProgramIROp_EntryPointOp, - RK_Op_ProgramIROp_AllocOp, - RK_Op_ProgramIROp_FreeOp, - RK_Op_ProgramIROp_ModeConfigOp, - RK_Op_ProgramIROp_BindOp, - RK_Op_ProgramIROp_Last, - RK_Op_DbgIROp, - RK_Op_DbgIROp_CommentOp, - RK_Op_DbgIROp_HintsOp, - RK_Op_DbgIROp_Last, - RK_Op_Last, - RK_Val, - RK_Val_LinalgIRVal, - RK_Val_GraphIRVal, - RK_Val_TensorIRVal, - RK_Val_TensorIRVal_TensorVal, - RK_Val_TensorIRVal_Last, - RK_Val_BuiltinIRVal, - RK_Val_ControlFlowIRVal, - RK_Val_ProgramIRVal, - RK_Val_DbgIRVal, - RK_Val_Last, - RK_Attr, - RK_Attr_LinalgIRAttr, - RK_Attr_LinalgIRAttr_QuantizationAnnotation, - RK_Attr_LinalgIRAttr_QuantizationSpec, - RK_Attr_LinalgIRAttr_Last, - RK_Attr_GraphIRAttr, - RK_Attr_TensorIRAttr, - RK_Attr_BuiltinIRAttr, - RK_Attr_BuiltinIRAttr_IntAttr, - RK_Attr_BuiltinIRAttr_FPAttr, - RK_Attr_BuiltinIRAttr_StrAttr, - RK_Attr_BuiltinIRAttr_SymbolAttr, - RK_Attr_BuiltinIRAttr_BoolAttr, - RK_Attr_BuiltinIRAttr_VectorFP32Attr, - RK_Attr_BuiltinIRAttr_VectorInt16Attr, - RK_Attr_BuiltinIRAttr_VectorUInt16Attr, - RK_Attr_BuiltinIRAttr_Last, - RK_Attr_ControlFlowIRAttr, - RK_Attr_ProgramIRAttr, - RK_Attr_DbgIRAttr, - RK_Attr_Last, +RK_None, +RK_Op, +RK_Op_LinalgIROp, +RK_Op_LinalgIROp_RegisterOp, +RK_Op_LinalgIROp_CustomKernelOp, +RK_Op_LinalgIROp_FillOp, +RK_Op_LinalgIROp_AddOp, +RK_Op_LinalgIROp_SubOp, +RK_Op_LinalgIROp_STFTOp, +RK_Op_LinalgIROp_ISTFTOp, +RK_Op_LinalgIROp_MulOp, +RK_Op_LinalgIROp_DivOp, +RK_Op_LinalgIROp_AbsOp, +RK_Op_LinalgIROp_LogOp, +RK_Op_LinalgIROp_MatMulOp, +RK_Op_LinalgIROp_EmbeddingOp, +RK_Op_LinalgIROp_LinearOp, +RK_Op_LinalgIROp_RoPEOp, +RK_Op_LinalgIROp_SoftmaxOp, +RK_Op_LinalgIROp_TransposeOp, +RK_Op_LinalgIROp_RMSNormOp, +RK_Op_LinalgIROp_SiLUOp, +RK_Op_LinalgIROp_KVCacheOp, +RK_Op_LinalgIROp_CausalMaskOp, +RK_Op_LinalgIROp_CastTypeOp, +RK_Op_LinalgIROp_X2XOp, +RK_Op_LinalgIROp_ViewOp, +RK_Op_LinalgIROp_SplitOp, +RK_Op_LinalgIROp_FlashAttention2Op, +RK_Op_LinalgIROp_CausalDepthwiseConv1DOp, +RK_Op_LinalgIROp_GroupedQueryAttentionOp, +RK_Op_LinalgIROp_ParallelLinearOp, +RK_Op_LinalgIROp_GatedDeltaRuleOp, +RK_Op_LinalgIROp_RepeatOp, +RK_Op_LinalgIROp_PermuteOp, +RK_Op_LinalgIROp_Conv1DOp, +RK_Op_LinalgIROp_Conv2DOp, +RK_Op_LinalgIROp_Conv3DOp, +RK_Op_LinalgIROp_GELUOp, +RK_Op_LinalgIROp_LayerNormOp, +RK_Op_LinalgIROp_MultimodalRoPEOp, +RK_Op_LinalgIROp_VisionRoPEOp, +RK_Op_LinalgIROp_QuickGELUOp, +RK_Op_LinalgIROp_KimiDeltaAttentionOp, +RK_Op_LinalgIROp_CopyOp, +RK_Op_LinalgIROp_CloneOp, +RK_Op_LinalgIROp_NegOp, +RK_Op_LinalgIROp_ConcatOp, +RK_Op_LinalgIROp_ReduceMinOp, +RK_Op_LinalgIROp_ReduceMaxOp, +RK_Op_LinalgIROp_ReduceSumOp, +RK_Op_LinalgIROp_ReLUOp, +RK_Op_LinalgIROp_ContiguousOp, +RK_Op_LinalgIROp_ReshapeOp, +RK_Op_LinalgIROp_SliceOp, +RK_Op_LinalgIROp_ParamOp, +RK_Op_LinalgIROp_IndexOp, +RK_Op_LinalgIROp_TopKOp, +RK_Op_LinalgIROp_MeanOp, +RK_Op_LinalgIROp_ClipOp, +RK_Op_LinalgIROp_ExpOp, +RK_Op_LinalgIROp_SinOp, +RK_Op_LinalgIROp_CosOp, +RK_Op_LinalgIROp_PagedAttnOp, +RK_Op_LinalgIROp_LayerNorm2DOp, +RK_Op_LinalgIROp_PadOp, +RK_Op_LinalgIROp_InterpolateOp, +RK_Op_LinalgIROp_EinsumOp, +RK_Op_LinalgIROp_StackOp, +RK_Op_LinalgIROp_MaskedScatterOp, +RK_Op_LinalgIROp_ScatterOp, +RK_Op_LinalgIROp_GatherOp, +RK_Op_LinalgIROp_ArgsortOp, +RK_Op_LinalgIROp_FlashAttention2SwaSinkOp, +RK_Op_LinalgIROp_RadixAttnRelaxOp, +RK_Op_LinalgIROp_RadixAttnSwaSinkOp, +RK_Op_LinalgIROp_EqualOp, +RK_Op_LinalgIROp_WhereOp, +RK_Op_LinalgIROp_SigmoidOp, +RK_Op_LinalgIROp_CustomizedOp, +RK_Op_LinalgIROp_Last, +RK_Op_GraphIROp, +RK_Op_GraphIROp_SubGraphOp, +RK_Op_GraphIROp_CallGraphOp, +RK_Op_GraphIROp_Last, +RK_Op_TensorIROp, +RK_Op_TensorIROp_AllocOp, +RK_Op_TensorIROp_RegisterOp, +RK_Op_TensorIROp_FreeOp, +RK_Op_TensorIROp_Last, +RK_Op_BuiltinIROp, +RK_Op_BuiltinIROp_ModuleOp, +RK_Op_BuiltinIROp_Last, +RK_Op_ControlFlowIROp, +RK_Op_ControlFlowIROp_ReturnOp, +RK_Op_ControlFlowIROp_Last, +RK_Op_ProgramIROp, +RK_Op_ProgramIROp_FragmentOp, +RK_Op_ProgramIROp_KernelLaunchOp, +RK_Op_ProgramIROp_KernelSymbolOp, +RK_Op_ProgramIROp_ValueSymbolOp, +RK_Op_ProgramIROp_JumpOp, +RK_Op_ProgramIROp_LabelOp, +RK_Op_ProgramIROp_ExitOp, +RK_Op_ProgramIROp_RetOp, +RK_Op_ProgramIROp_EntryPointOp, +RK_Op_ProgramIROp_AllocOp, +RK_Op_ProgramIROp_FreeOp, +RK_Op_ProgramIROp_ModeConfigOp, +RK_Op_ProgramIROp_BindOp, +RK_Op_ProgramIROp_Last, +RK_Op_DbgIROp, +RK_Op_DbgIROp_CommentOp, +RK_Op_DbgIROp_HintsOp, +RK_Op_DbgIROp_Last, +RK_Op_Last, +RK_Val, +RK_Val_LinalgIRVal, +RK_Val_GraphIRVal, +RK_Val_TensorIRVal, +RK_Val_TensorIRVal_TensorVal, +RK_Val_TensorIRVal_Last, +RK_Val_BuiltinIRVal, +RK_Val_ControlFlowIRVal, +RK_Val_ProgramIRVal, +RK_Val_DbgIRVal, +RK_Val_Last, +RK_Attr, +RK_Attr_LinalgIRAttr, +RK_Attr_LinalgIRAttr_QuantizationAnnotation, +RK_Attr_LinalgIRAttr_QuantizationSpec, +RK_Attr_LinalgIRAttr_Last, +RK_Attr_GraphIRAttr, +RK_Attr_TensorIRAttr, +RK_Attr_BuiltinIRAttr, +RK_Attr_BuiltinIRAttr_IntAttr, +RK_Attr_BuiltinIRAttr_FPAttr, +RK_Attr_BuiltinIRAttr_StrAttr, +RK_Attr_BuiltinIRAttr_SymbolAttr, +RK_Attr_BuiltinIRAttr_BoolAttr, +RK_Attr_BuiltinIRAttr_VectorFP32Attr, +RK_Attr_BuiltinIRAttr_VectorInt16Attr, +RK_Attr_BuiltinIRAttr_VectorUInt16Attr, +RK_Attr_BuiltinIRAttr_Last, +RK_Attr_ControlFlowIRAttr, +RK_Attr_ProgramIRAttr, +RK_Attr_DbgIRAttr, +RK_Attr_Last, }; + class Node; template struct NodeRTTIClassOfImpl { static inline bool classof(Node* v) { return false; } }; -} // namespace mllm::ir +} + \ No newline at end of file diff --git a/mllm/compile/ir/NodeRTTIClassOfImpl.hpp b/mllm/compile/ir/NodeRTTIClassOfImpl.hpp index 252c3fe4d..bde2105f3 100644 --- a/mllm/compile/ir/NodeRTTIClassOfImpl.hpp +++ b/mllm/compile/ir/NodeRTTIClassOfImpl.hpp @@ -1,4 +1,4 @@ -// Auto generated: 2026-08-21 15:20:12 +// Auto generated: 2026-09-06 09:09:14 // do not modify this file #pragma once namespace mllm::ir { @@ -9,392 +9,408 @@ template struct NodeRTTIClassOfImpl { static inline bool classof(Node* v) { return false; } }; -#endif //! RTTI_Node_IMPL +#endif //! RTTI_Node_IMPL -#define RTTI_RK_OP_IMPL(v) return (v)->getKind() >= RK_Op && (v)->getKind() <= RK_Op_Last +#define RTTI_RK_OP_IMPL(v) \ + return (v)->getKind() >= RK_Op && (v)->getKind() <= RK_Op_Last \ -#define RTTI_RK_OP_LINALGIROP_IMPL(v) return (v)->getKind() >= RK_Op_LinalgIROp && (v)->getKind() <= RK_Op_LinalgIROp_Last +#define RTTI_RK_OP_LINALGIROP_IMPL(v) \ + return (v)->getKind() >= RK_Op_LinalgIROp && (v)->getKind() <= RK_Op_LinalgIROp_Last \ #define RTTI_RK_OP_LINALGIROP_REGISTEROP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_RegisterOp && (v)->getKind() <= RK_Op_LinalgIROp_RegisterOp + return (v)->getKind() >= RK_Op_LinalgIROp_RegisterOp && (v)->getKind() <= RK_Op_LinalgIROp_RegisterOp \ #define RTTI_RK_OP_LINALGIROP_CUSTOMKERNELOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_CustomKernelOp && (v)->getKind() <= RK_Op_LinalgIROp_CustomKernelOp + return (v)->getKind() >= RK_Op_LinalgIROp_CustomKernelOp && (v)->getKind() <= RK_Op_LinalgIROp_CustomKernelOp \ #define RTTI_RK_OP_LINALGIROP_FILLOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_FillOp && (v)->getKind() <= RK_Op_LinalgIROp_FillOp + return (v)->getKind() >= RK_Op_LinalgIROp_FillOp && (v)->getKind() <= RK_Op_LinalgIROp_FillOp \ #define RTTI_RK_OP_LINALGIROP_ADDOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_AddOp && (v)->getKind() <= RK_Op_LinalgIROp_AddOp + return (v)->getKind() >= RK_Op_LinalgIROp_AddOp && (v)->getKind() <= RK_Op_LinalgIROp_AddOp \ #define RTTI_RK_OP_LINALGIROP_SUBOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_SubOp && (v)->getKind() <= RK_Op_LinalgIROp_SubOp + return (v)->getKind() >= RK_Op_LinalgIROp_SubOp && (v)->getKind() <= RK_Op_LinalgIROp_SubOp \ #define RTTI_RK_OP_LINALGIROP_STFTOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_STFTOp && (v)->getKind() <= RK_Op_LinalgIROp_STFTOp + return (v)->getKind() >= RK_Op_LinalgIROp_STFTOp && (v)->getKind() <= RK_Op_LinalgIROp_STFTOp \ #define RTTI_RK_OP_LINALGIROP_ISTFTOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ISTFTOp && (v)->getKind() <= RK_Op_LinalgIROp_ISTFTOp + return (v)->getKind() >= RK_Op_LinalgIROp_ISTFTOp && (v)->getKind() <= RK_Op_LinalgIROp_ISTFTOp \ #define RTTI_RK_OP_LINALGIROP_MULOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_MulOp && (v)->getKind() <= RK_Op_LinalgIROp_MulOp + return (v)->getKind() >= RK_Op_LinalgIROp_MulOp && (v)->getKind() <= RK_Op_LinalgIROp_MulOp \ #define RTTI_RK_OP_LINALGIROP_DIVOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_DivOp && (v)->getKind() <= RK_Op_LinalgIROp_DivOp + return (v)->getKind() >= RK_Op_LinalgIROp_DivOp && (v)->getKind() <= RK_Op_LinalgIROp_DivOp \ #define RTTI_RK_OP_LINALGIROP_ABSOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_AbsOp && (v)->getKind() <= RK_Op_LinalgIROp_AbsOp + return (v)->getKind() >= RK_Op_LinalgIROp_AbsOp && (v)->getKind() <= RK_Op_LinalgIROp_AbsOp \ #define RTTI_RK_OP_LINALGIROP_LOGOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_LogOp && (v)->getKind() <= RK_Op_LinalgIROp_LogOp + return (v)->getKind() >= RK_Op_LinalgIROp_LogOp && (v)->getKind() <= RK_Op_LinalgIROp_LogOp \ #define RTTI_RK_OP_LINALGIROP_MATMULOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_MatMulOp && (v)->getKind() <= RK_Op_LinalgIROp_MatMulOp + return (v)->getKind() >= RK_Op_LinalgIROp_MatMulOp && (v)->getKind() <= RK_Op_LinalgIROp_MatMulOp \ #define RTTI_RK_OP_LINALGIROP_EMBEDDINGOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_EmbeddingOp && (v)->getKind() <= RK_Op_LinalgIROp_EmbeddingOp + return (v)->getKind() >= RK_Op_LinalgIROp_EmbeddingOp && (v)->getKind() <= RK_Op_LinalgIROp_EmbeddingOp \ #define RTTI_RK_OP_LINALGIROP_LINEAROP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_LinearOp && (v)->getKind() <= RK_Op_LinalgIROp_LinearOp + return (v)->getKind() >= RK_Op_LinalgIROp_LinearOp && (v)->getKind() <= RK_Op_LinalgIROp_LinearOp \ #define RTTI_RK_OP_LINALGIROP_ROPEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_RoPEOp && (v)->getKind() <= RK_Op_LinalgIROp_RoPEOp + return (v)->getKind() >= RK_Op_LinalgIROp_RoPEOp && (v)->getKind() <= RK_Op_LinalgIROp_RoPEOp \ #define RTTI_RK_OP_LINALGIROP_SOFTMAXOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_SoftmaxOp && (v)->getKind() <= RK_Op_LinalgIROp_SoftmaxOp + return (v)->getKind() >= RK_Op_LinalgIROp_SoftmaxOp && (v)->getKind() <= RK_Op_LinalgIROp_SoftmaxOp \ #define RTTI_RK_OP_LINALGIROP_TRANSPOSEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_TransposeOp && (v)->getKind() <= RK_Op_LinalgIROp_TransposeOp + return (v)->getKind() >= RK_Op_LinalgIROp_TransposeOp && (v)->getKind() <= RK_Op_LinalgIROp_TransposeOp \ #define RTTI_RK_OP_LINALGIROP_RMSNORMOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_RMSNormOp && (v)->getKind() <= RK_Op_LinalgIROp_RMSNormOp + return (v)->getKind() >= RK_Op_LinalgIROp_RMSNormOp && (v)->getKind() <= RK_Op_LinalgIROp_RMSNormOp \ #define RTTI_RK_OP_LINALGIROP_SILUOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_SiLUOp && (v)->getKind() <= RK_Op_LinalgIROp_SiLUOp + return (v)->getKind() >= RK_Op_LinalgIROp_SiLUOp && (v)->getKind() <= RK_Op_LinalgIROp_SiLUOp \ #define RTTI_RK_OP_LINALGIROP_KVCACHEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_KVCacheOp && (v)->getKind() <= RK_Op_LinalgIROp_KVCacheOp + return (v)->getKind() >= RK_Op_LinalgIROp_KVCacheOp && (v)->getKind() <= RK_Op_LinalgIROp_KVCacheOp \ #define RTTI_RK_OP_LINALGIROP_CAUSALMASKOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_CausalMaskOp && (v)->getKind() <= RK_Op_LinalgIROp_CausalMaskOp + return (v)->getKind() >= RK_Op_LinalgIROp_CausalMaskOp && (v)->getKind() <= RK_Op_LinalgIROp_CausalMaskOp \ #define RTTI_RK_OP_LINALGIROP_CASTTYPEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_CastTypeOp && (v)->getKind() <= RK_Op_LinalgIROp_CastTypeOp + return (v)->getKind() >= RK_Op_LinalgIROp_CastTypeOp && (v)->getKind() <= RK_Op_LinalgIROp_CastTypeOp \ #define RTTI_RK_OP_LINALGIROP_X2XOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_X2XOp && (v)->getKind() <= RK_Op_LinalgIROp_X2XOp + return (v)->getKind() >= RK_Op_LinalgIROp_X2XOp && (v)->getKind() <= RK_Op_LinalgIROp_X2XOp \ #define RTTI_RK_OP_LINALGIROP_VIEWOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ViewOp && (v)->getKind() <= RK_Op_LinalgIROp_ViewOp + return (v)->getKind() >= RK_Op_LinalgIROp_ViewOp && (v)->getKind() <= RK_Op_LinalgIROp_ViewOp \ #define RTTI_RK_OP_LINALGIROP_SPLITOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_SplitOp && (v)->getKind() <= RK_Op_LinalgIROp_SplitOp + return (v)->getKind() >= RK_Op_LinalgIROp_SplitOp && (v)->getKind() <= RK_Op_LinalgIROp_SplitOp \ #define RTTI_RK_OP_LINALGIROP_FLASHATTENTION2OP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_FlashAttention2Op && (v)->getKind() <= RK_Op_LinalgIROp_FlashAttention2Op + return (v)->getKind() >= RK_Op_LinalgIROp_FlashAttention2Op && (v)->getKind() <= RK_Op_LinalgIROp_FlashAttention2Op \ -#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_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_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 + return (v)->getKind() >= RK_Op_LinalgIROp_ParallelLinearOp && (v)->getKind() <= RK_Op_LinalgIROp_ParallelLinearOp \ #define RTTI_RK_OP_LINALGIROP_GATEDDELTARULEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_GatedDeltaRuleOp && (v)->getKind() <= RK_Op_LinalgIROp_GatedDeltaRuleOp + return (v)->getKind() >= RK_Op_LinalgIROp_GatedDeltaRuleOp && (v)->getKind() <= RK_Op_LinalgIROp_GatedDeltaRuleOp \ #define RTTI_RK_OP_LINALGIROP_REPEATOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_RepeatOp && (v)->getKind() <= RK_Op_LinalgIROp_RepeatOp + return (v)->getKind() >= RK_Op_LinalgIROp_RepeatOp && (v)->getKind() <= RK_Op_LinalgIROp_RepeatOp \ #define RTTI_RK_OP_LINALGIROP_PERMUTEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_PermuteOp && (v)->getKind() <= RK_Op_LinalgIROp_PermuteOp + return (v)->getKind() >= RK_Op_LinalgIROp_PermuteOp && (v)->getKind() <= RK_Op_LinalgIROp_PermuteOp \ #define RTTI_RK_OP_LINALGIROP_CONV1DOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_Conv1DOp && (v)->getKind() <= RK_Op_LinalgIROp_Conv1DOp + return (v)->getKind() >= RK_Op_LinalgIROp_Conv1DOp && (v)->getKind() <= RK_Op_LinalgIROp_Conv1DOp \ #define RTTI_RK_OP_LINALGIROP_CONV2DOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_Conv2DOp && (v)->getKind() <= RK_Op_LinalgIROp_Conv2DOp + return (v)->getKind() >= RK_Op_LinalgIROp_Conv2DOp && (v)->getKind() <= RK_Op_LinalgIROp_Conv2DOp \ #define RTTI_RK_OP_LINALGIROP_CONV3DOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_Conv3DOp && (v)->getKind() <= RK_Op_LinalgIROp_Conv3DOp + return (v)->getKind() >= RK_Op_LinalgIROp_Conv3DOp && (v)->getKind() <= RK_Op_LinalgIROp_Conv3DOp \ #define RTTI_RK_OP_LINALGIROP_GELUOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_GELUOp && (v)->getKind() <= RK_Op_LinalgIROp_GELUOp + return (v)->getKind() >= RK_Op_LinalgIROp_GELUOp && (v)->getKind() <= RK_Op_LinalgIROp_GELUOp \ #define RTTI_RK_OP_LINALGIROP_LAYERNORMOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_LayerNormOp && (v)->getKind() <= RK_Op_LinalgIROp_LayerNormOp + return (v)->getKind() >= RK_Op_LinalgIROp_LayerNormOp && (v)->getKind() <= RK_Op_LinalgIROp_LayerNormOp \ #define RTTI_RK_OP_LINALGIROP_MULTIMODALROPEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_MultimodalRoPEOp && (v)->getKind() <= RK_Op_LinalgIROp_MultimodalRoPEOp + return (v)->getKind() >= RK_Op_LinalgIROp_MultimodalRoPEOp && (v)->getKind() <= RK_Op_LinalgIROp_MultimodalRoPEOp \ #define RTTI_RK_OP_LINALGIROP_VISIONROPEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_VisionRoPEOp && (v)->getKind() <= RK_Op_LinalgIROp_VisionRoPEOp + return (v)->getKind() >= RK_Op_LinalgIROp_VisionRoPEOp && (v)->getKind() <= RK_Op_LinalgIROp_VisionRoPEOp \ #define RTTI_RK_OP_LINALGIROP_QUICKGELUOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_QuickGELUOp && (v)->getKind() <= RK_Op_LinalgIROp_QuickGELUOp + return (v)->getKind() >= RK_Op_LinalgIROp_QuickGELUOp && (v)->getKind() <= RK_Op_LinalgIROp_QuickGELUOp \ + +#define RTTI_RK_OP_LINALGIROP_KIMIDELTAATTENTIONOP_IMPL(v) \ + return (v)->getKind() >= RK_Op_LinalgIROp_KimiDeltaAttentionOp && (v)->getKind() <= RK_Op_LinalgIROp_KimiDeltaAttentionOp \ #define RTTI_RK_OP_LINALGIROP_COPYOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_CopyOp && (v)->getKind() <= RK_Op_LinalgIROp_CopyOp + return (v)->getKind() >= RK_Op_LinalgIROp_CopyOp && (v)->getKind() <= RK_Op_LinalgIROp_CopyOp \ #define RTTI_RK_OP_LINALGIROP_CLONEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_CloneOp && (v)->getKind() <= RK_Op_LinalgIROp_CloneOp + return (v)->getKind() >= RK_Op_LinalgIROp_CloneOp && (v)->getKind() <= RK_Op_LinalgIROp_CloneOp \ #define RTTI_RK_OP_LINALGIROP_NEGOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_NegOp && (v)->getKind() <= RK_Op_LinalgIROp_NegOp + return (v)->getKind() >= RK_Op_LinalgIROp_NegOp && (v)->getKind() <= RK_Op_LinalgIROp_NegOp \ #define RTTI_RK_OP_LINALGIROP_CONCATOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ConcatOp && (v)->getKind() <= RK_Op_LinalgIROp_ConcatOp + return (v)->getKind() >= RK_Op_LinalgIROp_ConcatOp && (v)->getKind() <= RK_Op_LinalgIROp_ConcatOp \ #define RTTI_RK_OP_LINALGIROP_REDUCEMINOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ReduceMinOp && (v)->getKind() <= RK_Op_LinalgIROp_ReduceMinOp + return (v)->getKind() >= RK_Op_LinalgIROp_ReduceMinOp && (v)->getKind() <= RK_Op_LinalgIROp_ReduceMinOp \ #define RTTI_RK_OP_LINALGIROP_REDUCEMAXOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ReduceMaxOp && (v)->getKind() <= RK_Op_LinalgIROp_ReduceMaxOp + return (v)->getKind() >= RK_Op_LinalgIROp_ReduceMaxOp && (v)->getKind() <= RK_Op_LinalgIROp_ReduceMaxOp \ #define RTTI_RK_OP_LINALGIROP_REDUCESUMOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ReduceSumOp && (v)->getKind() <= RK_Op_LinalgIROp_ReduceSumOp + return (v)->getKind() >= RK_Op_LinalgIROp_ReduceSumOp && (v)->getKind() <= RK_Op_LinalgIROp_ReduceSumOp \ #define RTTI_RK_OP_LINALGIROP_RELUOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ReLUOp && (v)->getKind() <= RK_Op_LinalgIROp_ReLUOp + return (v)->getKind() >= RK_Op_LinalgIROp_ReLUOp && (v)->getKind() <= RK_Op_LinalgIROp_ReLUOp \ #define RTTI_RK_OP_LINALGIROP_CONTIGUOUSOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ContiguousOp && (v)->getKind() <= RK_Op_LinalgIROp_ContiguousOp + return (v)->getKind() >= RK_Op_LinalgIROp_ContiguousOp && (v)->getKind() <= RK_Op_LinalgIROp_ContiguousOp \ #define RTTI_RK_OP_LINALGIROP_RESHAPEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ReshapeOp && (v)->getKind() <= RK_Op_LinalgIROp_ReshapeOp + return (v)->getKind() >= RK_Op_LinalgIROp_ReshapeOp && (v)->getKind() <= RK_Op_LinalgIROp_ReshapeOp \ #define RTTI_RK_OP_LINALGIROP_SLICEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_SliceOp && (v)->getKind() <= RK_Op_LinalgIROp_SliceOp + return (v)->getKind() >= RK_Op_LinalgIROp_SliceOp && (v)->getKind() <= RK_Op_LinalgIROp_SliceOp \ #define RTTI_RK_OP_LINALGIROP_PARAMOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ParamOp && (v)->getKind() <= RK_Op_LinalgIROp_ParamOp + return (v)->getKind() >= RK_Op_LinalgIROp_ParamOp && (v)->getKind() <= RK_Op_LinalgIROp_ParamOp \ #define RTTI_RK_OP_LINALGIROP_INDEXOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_IndexOp && (v)->getKind() <= RK_Op_LinalgIROp_IndexOp + return (v)->getKind() >= RK_Op_LinalgIROp_IndexOp && (v)->getKind() <= RK_Op_LinalgIROp_IndexOp \ #define RTTI_RK_OP_LINALGIROP_TOPKOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_TopKOp && (v)->getKind() <= RK_Op_LinalgIROp_TopKOp + return (v)->getKind() >= RK_Op_LinalgIROp_TopKOp && (v)->getKind() <= RK_Op_LinalgIROp_TopKOp \ #define RTTI_RK_OP_LINALGIROP_MEANOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_MeanOp && (v)->getKind() <= RK_Op_LinalgIROp_MeanOp + return (v)->getKind() >= RK_Op_LinalgIROp_MeanOp && (v)->getKind() <= RK_Op_LinalgIROp_MeanOp \ #define RTTI_RK_OP_LINALGIROP_CLIPOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ClipOp && (v)->getKind() <= RK_Op_LinalgIROp_ClipOp + return (v)->getKind() >= RK_Op_LinalgIROp_ClipOp && (v)->getKind() <= RK_Op_LinalgIROp_ClipOp \ #define RTTI_RK_OP_LINALGIROP_EXPOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ExpOp && (v)->getKind() <= RK_Op_LinalgIROp_ExpOp + return (v)->getKind() >= RK_Op_LinalgIROp_ExpOp && (v)->getKind() <= RK_Op_LinalgIROp_ExpOp \ #define RTTI_RK_OP_LINALGIROP_SINOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_SinOp && (v)->getKind() <= RK_Op_LinalgIROp_SinOp + return (v)->getKind() >= RK_Op_LinalgIROp_SinOp && (v)->getKind() <= RK_Op_LinalgIROp_SinOp \ #define RTTI_RK_OP_LINALGIROP_COSOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_CosOp && (v)->getKind() <= RK_Op_LinalgIROp_CosOp + return (v)->getKind() >= RK_Op_LinalgIROp_CosOp && (v)->getKind() <= RK_Op_LinalgIROp_CosOp \ #define RTTI_RK_OP_LINALGIROP_PAGEDATTNOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_PagedAttnOp && (v)->getKind() <= RK_Op_LinalgIROp_PagedAttnOp + return (v)->getKind() >= RK_Op_LinalgIROp_PagedAttnOp && (v)->getKind() <= RK_Op_LinalgIROp_PagedAttnOp \ #define RTTI_RK_OP_LINALGIROP_LAYERNORM2DOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_LayerNorm2DOp && (v)->getKind() <= RK_Op_LinalgIROp_LayerNorm2DOp + return (v)->getKind() >= RK_Op_LinalgIROp_LayerNorm2DOp && (v)->getKind() <= RK_Op_LinalgIROp_LayerNorm2DOp \ #define RTTI_RK_OP_LINALGIROP_PADOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_PadOp && (v)->getKind() <= RK_Op_LinalgIROp_PadOp + return (v)->getKind() >= RK_Op_LinalgIROp_PadOp && (v)->getKind() <= RK_Op_LinalgIROp_PadOp \ #define RTTI_RK_OP_LINALGIROP_INTERPOLATEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_InterpolateOp && (v)->getKind() <= RK_Op_LinalgIROp_InterpolateOp + return (v)->getKind() >= RK_Op_LinalgIROp_InterpolateOp && (v)->getKind() <= RK_Op_LinalgIROp_InterpolateOp \ #define RTTI_RK_OP_LINALGIROP_EINSUMOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_EinsumOp && (v)->getKind() <= RK_Op_LinalgIROp_EinsumOp + return (v)->getKind() >= RK_Op_LinalgIROp_EinsumOp && (v)->getKind() <= RK_Op_LinalgIROp_EinsumOp \ #define RTTI_RK_OP_LINALGIROP_STACKOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_StackOp && (v)->getKind() <= RK_Op_LinalgIROp_StackOp + return (v)->getKind() >= RK_Op_LinalgIROp_StackOp && (v)->getKind() <= RK_Op_LinalgIROp_StackOp \ #define RTTI_RK_OP_LINALGIROP_MASKEDSCATTEROP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_MaskedScatterOp && (v)->getKind() <= RK_Op_LinalgIROp_MaskedScatterOp + return (v)->getKind() >= RK_Op_LinalgIROp_MaskedScatterOp && (v)->getKind() <= RK_Op_LinalgIROp_MaskedScatterOp \ #define RTTI_RK_OP_LINALGIROP_SCATTEROP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ScatterOp && (v)->getKind() <= RK_Op_LinalgIROp_ScatterOp + return (v)->getKind() >= RK_Op_LinalgIROp_ScatterOp && (v)->getKind() <= RK_Op_LinalgIROp_ScatterOp \ #define RTTI_RK_OP_LINALGIROP_GATHEROP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_GatherOp && (v)->getKind() <= RK_Op_LinalgIROp_GatherOp + return (v)->getKind() >= RK_Op_LinalgIROp_GatherOp && (v)->getKind() <= RK_Op_LinalgIROp_GatherOp \ #define RTTI_RK_OP_LINALGIROP_ARGSORTOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_ArgsortOp && (v)->getKind() <= RK_Op_LinalgIROp_ArgsortOp + return (v)->getKind() >= RK_Op_LinalgIROp_ArgsortOp && (v)->getKind() <= RK_Op_LinalgIROp_ArgsortOp \ -#define RTTI_RK_OP_LINALGIROP_FLASHATTENTION2SWASINKOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_FlashAttention2SwaSinkOp \ - && (v)->getKind() <= RK_Op_LinalgIROp_FlashAttention2SwaSinkOp +#define RTTI_RK_OP_LINALGIROP_FLASHATTENTION2SWASINKOP_IMPL(v) \ + return (v)->getKind() >= RK_Op_LinalgIROp_FlashAttention2SwaSinkOp && (v)->getKind() <= RK_Op_LinalgIROp_FlashAttention2SwaSinkOp \ #define RTTI_RK_OP_LINALGIROP_RADIXATTNRELAXOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_RadixAttnRelaxOp && (v)->getKind() <= RK_Op_LinalgIROp_RadixAttnRelaxOp + return (v)->getKind() >= RK_Op_LinalgIROp_RadixAttnRelaxOp && (v)->getKind() <= RK_Op_LinalgIROp_RadixAttnRelaxOp \ #define RTTI_RK_OP_LINALGIROP_RADIXATTNSWASINKOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_RadixAttnSwaSinkOp && (v)->getKind() <= RK_Op_LinalgIROp_RadixAttnSwaSinkOp + return (v)->getKind() >= RK_Op_LinalgIROp_RadixAttnSwaSinkOp && (v)->getKind() <= RK_Op_LinalgIROp_RadixAttnSwaSinkOp \ #define RTTI_RK_OP_LINALGIROP_EQUALOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_EqualOp && (v)->getKind() <= RK_Op_LinalgIROp_EqualOp + return (v)->getKind() >= RK_Op_LinalgIROp_EqualOp && (v)->getKind() <= RK_Op_LinalgIROp_EqualOp \ #define RTTI_RK_OP_LINALGIROP_WHEREOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_WhereOp && (v)->getKind() <= RK_Op_LinalgIROp_WhereOp + return (v)->getKind() >= RK_Op_LinalgIROp_WhereOp && (v)->getKind() <= RK_Op_LinalgIROp_WhereOp \ #define RTTI_RK_OP_LINALGIROP_SIGMOIDOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_SigmoidOp && (v)->getKind() <= RK_Op_LinalgIROp_SigmoidOp + return (v)->getKind() >= RK_Op_LinalgIROp_SigmoidOp && (v)->getKind() <= RK_Op_LinalgIROp_SigmoidOp \ #define RTTI_RK_OP_LINALGIROP_CUSTOMIZEDOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_CustomizedOp && (v)->getKind() <= RK_Op_LinalgIROp_CustomizedOp + return (v)->getKind() >= RK_Op_LinalgIROp_CustomizedOp && (v)->getKind() <= RK_Op_LinalgIROp_CustomizedOp \ -#define RTTI_RK_OP_GRAPHIROP_IMPL(v) return (v)->getKind() >= RK_Op_GraphIROp && (v)->getKind() <= RK_Op_GraphIROp_Last +#define RTTI_RK_OP_GRAPHIROP_IMPL(v) \ + return (v)->getKind() >= RK_Op_GraphIROp && (v)->getKind() <= RK_Op_GraphIROp_Last \ #define RTTI_RK_OP_GRAPHIROP_SUBGRAPHOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_GraphIROp_SubGraphOp && (v)->getKind() <= RK_Op_GraphIROp_SubGraphOp + return (v)->getKind() >= RK_Op_GraphIROp_SubGraphOp && (v)->getKind() <= RK_Op_GraphIROp_SubGraphOp \ #define RTTI_RK_OP_GRAPHIROP_CALLGRAPHOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_GraphIROp_CallGraphOp && (v)->getKind() <= RK_Op_GraphIROp_CallGraphOp + return (v)->getKind() >= RK_Op_GraphIROp_CallGraphOp && (v)->getKind() <= RK_Op_GraphIROp_CallGraphOp \ -#define RTTI_RK_OP_TENSORIROP_IMPL(v) return (v)->getKind() >= RK_Op_TensorIROp && (v)->getKind() <= RK_Op_TensorIROp_Last +#define RTTI_RK_OP_TENSORIROP_IMPL(v) \ + return (v)->getKind() >= RK_Op_TensorIROp && (v)->getKind() <= RK_Op_TensorIROp_Last \ #define RTTI_RK_OP_TENSORIROP_ALLOCOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_TensorIROp_AllocOp && (v)->getKind() <= RK_Op_TensorIROp_AllocOp + return (v)->getKind() >= RK_Op_TensorIROp_AllocOp && (v)->getKind() <= RK_Op_TensorIROp_AllocOp \ #define RTTI_RK_OP_TENSORIROP_REGISTEROP_IMPL(v) \ - return (v)->getKind() >= RK_Op_TensorIROp_RegisterOp && (v)->getKind() <= RK_Op_TensorIROp_RegisterOp + return (v)->getKind() >= RK_Op_TensorIROp_RegisterOp && (v)->getKind() <= RK_Op_TensorIROp_RegisterOp \ #define RTTI_RK_OP_TENSORIROP_FREEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_TensorIROp_FreeOp && (v)->getKind() <= RK_Op_TensorIROp_FreeOp + return (v)->getKind() >= RK_Op_TensorIROp_FreeOp && (v)->getKind() <= RK_Op_TensorIROp_FreeOp \ -#define RTTI_RK_OP_BUILTINIROP_IMPL(v) return (v)->getKind() >= RK_Op_BuiltinIROp && (v)->getKind() <= RK_Op_BuiltinIROp_Last +#define RTTI_RK_OP_BUILTINIROP_IMPL(v) \ + return (v)->getKind() >= RK_Op_BuiltinIROp && (v)->getKind() <= RK_Op_BuiltinIROp_Last \ #define RTTI_RK_OP_BUILTINIROP_MODULEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_BuiltinIROp_ModuleOp && (v)->getKind() <= RK_Op_BuiltinIROp_ModuleOp + return (v)->getKind() >= RK_Op_BuiltinIROp_ModuleOp && (v)->getKind() <= RK_Op_BuiltinIROp_ModuleOp \ #define RTTI_RK_OP_CONTROLFLOWIROP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ControlFlowIROp && (v)->getKind() <= RK_Op_ControlFlowIROp_Last + return (v)->getKind() >= RK_Op_ControlFlowIROp && (v)->getKind() <= RK_Op_ControlFlowIROp_Last \ #define RTTI_RK_OP_CONTROLFLOWIROP_RETURNOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ControlFlowIROp_ReturnOp && (v)->getKind() <= RK_Op_ControlFlowIROp_ReturnOp + return (v)->getKind() >= RK_Op_ControlFlowIROp_ReturnOp && (v)->getKind() <= RK_Op_ControlFlowIROp_ReturnOp \ -#define RTTI_RK_OP_PROGRAMIROP_IMPL(v) return (v)->getKind() >= RK_Op_ProgramIROp && (v)->getKind() <= RK_Op_ProgramIROp_Last +#define RTTI_RK_OP_PROGRAMIROP_IMPL(v) \ + return (v)->getKind() >= RK_Op_ProgramIROp && (v)->getKind() <= RK_Op_ProgramIROp_Last \ #define RTTI_RK_OP_PROGRAMIROP_FRAGMENTOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_FragmentOp && (v)->getKind() <= RK_Op_ProgramIROp_FragmentOp + return (v)->getKind() >= RK_Op_ProgramIROp_FragmentOp && (v)->getKind() <= RK_Op_ProgramIROp_FragmentOp \ #define RTTI_RK_OP_PROGRAMIROP_KERNELLAUNCHOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_KernelLaunchOp && (v)->getKind() <= RK_Op_ProgramIROp_KernelLaunchOp + return (v)->getKind() >= RK_Op_ProgramIROp_KernelLaunchOp && (v)->getKind() <= RK_Op_ProgramIROp_KernelLaunchOp \ #define RTTI_RK_OP_PROGRAMIROP_KERNELSYMBOLOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_KernelSymbolOp && (v)->getKind() <= RK_Op_ProgramIROp_KernelSymbolOp + return (v)->getKind() >= RK_Op_ProgramIROp_KernelSymbolOp && (v)->getKind() <= RK_Op_ProgramIROp_KernelSymbolOp \ #define RTTI_RK_OP_PROGRAMIROP_VALUESYMBOLOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_ValueSymbolOp && (v)->getKind() <= RK_Op_ProgramIROp_ValueSymbolOp + return (v)->getKind() >= RK_Op_ProgramIROp_ValueSymbolOp && (v)->getKind() <= RK_Op_ProgramIROp_ValueSymbolOp \ #define RTTI_RK_OP_PROGRAMIROP_JUMPOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_JumpOp && (v)->getKind() <= RK_Op_ProgramIROp_JumpOp + return (v)->getKind() >= RK_Op_ProgramIROp_JumpOp && (v)->getKind() <= RK_Op_ProgramIROp_JumpOp \ #define RTTI_RK_OP_PROGRAMIROP_LABELOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_LabelOp && (v)->getKind() <= RK_Op_ProgramIROp_LabelOp + return (v)->getKind() >= RK_Op_ProgramIROp_LabelOp && (v)->getKind() <= RK_Op_ProgramIROp_LabelOp \ #define RTTI_RK_OP_PROGRAMIROP_EXITOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_ExitOp && (v)->getKind() <= RK_Op_ProgramIROp_ExitOp + return (v)->getKind() >= RK_Op_ProgramIROp_ExitOp && (v)->getKind() <= RK_Op_ProgramIROp_ExitOp \ #define RTTI_RK_OP_PROGRAMIROP_RETOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_RetOp && (v)->getKind() <= RK_Op_ProgramIROp_RetOp + return (v)->getKind() >= RK_Op_ProgramIROp_RetOp && (v)->getKind() <= RK_Op_ProgramIROp_RetOp \ #define RTTI_RK_OP_PROGRAMIROP_ENTRYPOINTOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_EntryPointOp && (v)->getKind() <= RK_Op_ProgramIROp_EntryPointOp + return (v)->getKind() >= RK_Op_ProgramIROp_EntryPointOp && (v)->getKind() <= RK_Op_ProgramIROp_EntryPointOp \ #define RTTI_RK_OP_PROGRAMIROP_ALLOCOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_AllocOp && (v)->getKind() <= RK_Op_ProgramIROp_AllocOp + return (v)->getKind() >= RK_Op_ProgramIROp_AllocOp && (v)->getKind() <= RK_Op_ProgramIROp_AllocOp \ #define RTTI_RK_OP_PROGRAMIROP_FREEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_FreeOp && (v)->getKind() <= RK_Op_ProgramIROp_FreeOp + return (v)->getKind() >= RK_Op_ProgramIROp_FreeOp && (v)->getKind() <= RK_Op_ProgramIROp_FreeOp \ #define RTTI_RK_OP_PROGRAMIROP_MODECONFIGOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_ModeConfigOp && (v)->getKind() <= RK_Op_ProgramIROp_ModeConfigOp + return (v)->getKind() >= RK_Op_ProgramIROp_ModeConfigOp && (v)->getKind() <= RK_Op_ProgramIROp_ModeConfigOp \ #define RTTI_RK_OP_PROGRAMIROP_BINDOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_ProgramIROp_BindOp && (v)->getKind() <= RK_Op_ProgramIROp_BindOp + return (v)->getKind() >= RK_Op_ProgramIROp_BindOp && (v)->getKind() <= RK_Op_ProgramIROp_BindOp \ -#define RTTI_RK_OP_DBGIROP_IMPL(v) return (v)->getKind() >= RK_Op_DbgIROp && (v)->getKind() <= RK_Op_DbgIROp_Last +#define RTTI_RK_OP_DBGIROP_IMPL(v) \ + return (v)->getKind() >= RK_Op_DbgIROp && (v)->getKind() <= RK_Op_DbgIROp_Last \ #define RTTI_RK_OP_DBGIROP_COMMENTOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_DbgIROp_CommentOp && (v)->getKind() <= RK_Op_DbgIROp_CommentOp + return (v)->getKind() >= RK_Op_DbgIROp_CommentOp && (v)->getKind() <= RK_Op_DbgIROp_CommentOp \ #define RTTI_RK_OP_DBGIROP_HINTSOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_DbgIROp_HintsOp && (v)->getKind() <= RK_Op_DbgIROp_HintsOp + return (v)->getKind() >= RK_Op_DbgIROp_HintsOp && (v)->getKind() <= RK_Op_DbgIROp_HintsOp \ -#define RTTI_RK_VAL_IMPL(v) return (v)->getKind() >= RK_Val && (v)->getKind() <= RK_Val_Last +#define RTTI_RK_VAL_IMPL(v) \ + return (v)->getKind() >= RK_Val && (v)->getKind() <= RK_Val_Last \ -#define RTTI_RK_VAL_LINALGIRVAL_IMPL(v) return (v)->getKind() >= RK_Val_LinalgIRVal && (v)->getKind() <= RK_Val_LinalgIRVal +#define RTTI_RK_VAL_LINALGIRVAL_IMPL(v) \ + return (v)->getKind() >= RK_Val_LinalgIRVal && (v)->getKind() <= RK_Val_LinalgIRVal \ -#define RTTI_RK_VAL_GRAPHIRVAL_IMPL(v) return (v)->getKind() >= RK_Val_GraphIRVal && (v)->getKind() <= RK_Val_GraphIRVal +#define RTTI_RK_VAL_GRAPHIRVAL_IMPL(v) \ + return (v)->getKind() >= RK_Val_GraphIRVal && (v)->getKind() <= RK_Val_GraphIRVal \ -#define RTTI_RK_VAL_TENSORIRVAL_IMPL(v) return (v)->getKind() >= RK_Val_TensorIRVal && (v)->getKind() <= RK_Val_TensorIRVal_Last +#define RTTI_RK_VAL_TENSORIRVAL_IMPL(v) \ + return (v)->getKind() >= RK_Val_TensorIRVal && (v)->getKind() <= RK_Val_TensorIRVal_Last \ #define RTTI_RK_VAL_TENSORIRVAL_TENSORVAL_IMPL(v) \ - return (v)->getKind() >= RK_Val_TensorIRVal_TensorVal && (v)->getKind() <= RK_Val_TensorIRVal_TensorVal + return (v)->getKind() >= RK_Val_TensorIRVal_TensorVal && (v)->getKind() <= RK_Val_TensorIRVal_TensorVal \ -#define RTTI_RK_VAL_BUILTINIRVAL_IMPL(v) return (v)->getKind() >= RK_Val_BuiltinIRVal && (v)->getKind() <= RK_Val_BuiltinIRVal +#define RTTI_RK_VAL_BUILTINIRVAL_IMPL(v) \ + return (v)->getKind() >= RK_Val_BuiltinIRVal && (v)->getKind() <= RK_Val_BuiltinIRVal \ #define RTTI_RK_VAL_CONTROLFLOWIRVAL_IMPL(v) \ - return (v)->getKind() >= RK_Val_ControlFlowIRVal && (v)->getKind() <= RK_Val_ControlFlowIRVal + return (v)->getKind() >= RK_Val_ControlFlowIRVal && (v)->getKind() <= RK_Val_ControlFlowIRVal \ -#define RTTI_RK_VAL_PROGRAMIRVAL_IMPL(v) return (v)->getKind() >= RK_Val_ProgramIRVal && (v)->getKind() <= RK_Val_ProgramIRVal +#define RTTI_RK_VAL_PROGRAMIRVAL_IMPL(v) \ + return (v)->getKind() >= RK_Val_ProgramIRVal && (v)->getKind() <= RK_Val_ProgramIRVal \ -#define RTTI_RK_VAL_DBGIRVAL_IMPL(v) return (v)->getKind() >= RK_Val_DbgIRVal && (v)->getKind() <= RK_Val_DbgIRVal +#define RTTI_RK_VAL_DBGIRVAL_IMPL(v) \ + return (v)->getKind() >= RK_Val_DbgIRVal && (v)->getKind() <= RK_Val_DbgIRVal \ -#define RTTI_RK_ATTR_IMPL(v) return (v)->getKind() >= RK_Attr && (v)->getKind() <= RK_Attr_Last +#define RTTI_RK_ATTR_IMPL(v) \ + return (v)->getKind() >= RK_Attr && (v)->getKind() <= RK_Attr_Last \ #define RTTI_RK_ATTR_LINALGIRATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_LinalgIRAttr && (v)->getKind() <= RK_Attr_LinalgIRAttr_Last + return (v)->getKind() >= RK_Attr_LinalgIRAttr && (v)->getKind() <= RK_Attr_LinalgIRAttr_Last \ -#define RTTI_RK_ATTR_LINALGIRATTR_QUANTIZATIONANNOTATION_IMPL(v) \ - return (v)->getKind() >= RK_Attr_LinalgIRAttr_QuantizationAnnotation \ - && (v)->getKind() <= RK_Attr_LinalgIRAttr_QuantizationAnnotation +#define RTTI_RK_ATTR_LINALGIRATTR_QUANTIZATIONANNOTATION_IMPL(v) \ + return (v)->getKind() >= RK_Attr_LinalgIRAttr_QuantizationAnnotation && (v)->getKind() <= RK_Attr_LinalgIRAttr_QuantizationAnnotation \ #define RTTI_RK_ATTR_LINALGIRATTR_QUANTIZATIONSPEC_IMPL(v) \ - return (v)->getKind() >= RK_Attr_LinalgIRAttr_QuantizationSpec && (v)->getKind() <= RK_Attr_LinalgIRAttr_QuantizationSpec + return (v)->getKind() >= RK_Attr_LinalgIRAttr_QuantizationSpec && (v)->getKind() <= RK_Attr_LinalgIRAttr_QuantizationSpec \ -#define RTTI_RK_ATTR_GRAPHIRATTR_IMPL(v) return (v)->getKind() >= RK_Attr_GraphIRAttr && (v)->getKind() <= RK_Attr_GraphIRAttr +#define RTTI_RK_ATTR_GRAPHIRATTR_IMPL(v) \ + return (v)->getKind() >= RK_Attr_GraphIRAttr && (v)->getKind() <= RK_Attr_GraphIRAttr \ #define RTTI_RK_ATTR_TENSORIRATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_TensorIRAttr && (v)->getKind() <= RK_Attr_TensorIRAttr + return (v)->getKind() >= RK_Attr_TensorIRAttr && (v)->getKind() <= RK_Attr_TensorIRAttr \ #define RTTI_RK_ATTR_BUILTINIRATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_BuiltinIRAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_Last + return (v)->getKind() >= RK_Attr_BuiltinIRAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_Last \ #define RTTI_RK_ATTR_BUILTINIRATTR_INTATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_BuiltinIRAttr_IntAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_IntAttr + return (v)->getKind() >= RK_Attr_BuiltinIRAttr_IntAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_IntAttr \ #define RTTI_RK_ATTR_BUILTINIRATTR_FPATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_BuiltinIRAttr_FPAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_FPAttr + return (v)->getKind() >= RK_Attr_BuiltinIRAttr_FPAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_FPAttr \ #define RTTI_RK_ATTR_BUILTINIRATTR_STRATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_BuiltinIRAttr_StrAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_StrAttr + return (v)->getKind() >= RK_Attr_BuiltinIRAttr_StrAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_StrAttr \ #define RTTI_RK_ATTR_BUILTINIRATTR_SYMBOLATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_BuiltinIRAttr_SymbolAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_SymbolAttr + return (v)->getKind() >= RK_Attr_BuiltinIRAttr_SymbolAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_SymbolAttr \ #define RTTI_RK_ATTR_BUILTINIRATTR_BOOLATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_BuiltinIRAttr_BoolAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_BoolAttr + return (v)->getKind() >= RK_Attr_BuiltinIRAttr_BoolAttr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_BoolAttr \ #define RTTI_RK_ATTR_BUILTINIRATTR_VECTORFP32ATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_BuiltinIRAttr_VectorFP32Attr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_VectorFP32Attr + return (v)->getKind() >= RK_Attr_BuiltinIRAttr_VectorFP32Attr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_VectorFP32Attr \ #define RTTI_RK_ATTR_BUILTINIRATTR_VECTORINT16ATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_BuiltinIRAttr_VectorInt16Attr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_VectorInt16Attr + return (v)->getKind() >= RK_Attr_BuiltinIRAttr_VectorInt16Attr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_VectorInt16Attr \ #define RTTI_RK_ATTR_BUILTINIRATTR_VECTORUINT16ATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_BuiltinIRAttr_VectorUInt16Attr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_VectorUInt16Attr + return (v)->getKind() >= RK_Attr_BuiltinIRAttr_VectorUInt16Attr && (v)->getKind() <= RK_Attr_BuiltinIRAttr_VectorUInt16Attr \ #define RTTI_RK_ATTR_CONTROLFLOWIRATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_ControlFlowIRAttr && (v)->getKind() <= RK_Attr_ControlFlowIRAttr + return (v)->getKind() >= RK_Attr_ControlFlowIRAttr && (v)->getKind() <= RK_Attr_ControlFlowIRAttr \ #define RTTI_RK_ATTR_PROGRAMIRATTR_IMPL(v) \ - return (v)->getKind() >= RK_Attr_ProgramIRAttr && (v)->getKind() <= RK_Attr_ProgramIRAttr + return (v)->getKind() >= RK_Attr_ProgramIRAttr && (v)->getKind() <= RK_Attr_ProgramIRAttr \ -#define RTTI_RK_ATTR_DBGIRATTR_IMPL(v) return (v)->getKind() >= RK_Attr_DbgIRAttr && (v)->getKind() <= RK_Attr_DbgIRAttr +#define RTTI_RK_ATTR_DBGIRATTR_IMPL(v) \ + return (v)->getKind() >= RK_Attr_DbgIRAttr && (v)->getKind() <= RK_Attr_DbgIRAttr \ -} // namespace mllm::ir +} diff --git a/mllm/compile/ir/linalg/Op.cpp b/mllm/compile/ir/linalg/Op.cpp index 3b315e6f2..a5a8f1bff 100644 --- a/mllm/compile/ir/linalg/Op.cpp +++ b/mllm/compile/ir/linalg/Op.cpp @@ -84,6 +84,7 @@ LINALG_AOPS_DECL(OpTypes::kMultimodalRoPE, MultimodalRoPEOp); LINALG_AOPS_DECL(OpTypes::kVisionRoPE, VisionRoPEOp); LINALG_AOPS_DECL(OpTypes::kQuickGELU, QuickGELUOp); +LINALG_AOPS_DECL(OpTypes::kKimiDeltaAttention, KimiDeltaAttentionOp); LINALG_AOPS_DECL(OpTypes::kCopy, CopyOp); LINALG_AOPS_DECL(OpTypes::kClone, CloneOp); diff --git a/mllm/compile/ir/linalg/Op.hpp b/mllm/compile/ir/linalg/Op.hpp index f95911cc9..d3514bcae 100644 --- a/mllm/compile/ir/linalg/Op.hpp +++ b/mllm/compile/ir/linalg/Op.hpp @@ -49,6 +49,7 @@ class LayerNormOp; class MultimodalRoPEOp; class VisionRoPEOp; class QuickGELUOp; +class KimiDeltaAttentionOp; class CopyOp; class CloneOp; class NegOp; @@ -219,6 +220,7 @@ LINALG_AOPS_DEFINE(MultimodalRoPEOp, MULTIMODALROPEOP); LINALG_AOPS_DEFINE(VisionRoPEOp, VISIONROPEOP); LINALG_AOPS_DEFINE(QuickGELUOp, QUICKGELUOP); +LINALG_AOPS_DEFINE(KimiDeltaAttentionOp, KIMIDELTAATTENTIONOP); LINALG_AOPS_DEFINE(CopyOp, COPYOP); LINALG_AOPS_DEFINE(CloneOp, CLONEOP); diff --git a/mllm/compile/ir/rtti_kind_gen.py b/mllm/compile/ir/rtti_kind_gen.py index eda139259..41dd654e5 100644 --- a/mllm/compile/ir/rtti_kind_gen.py +++ b/mllm/compile/ir/rtti_kind_gen.py @@ -261,6 +261,7 @@ def define_lianlg_ir(ir: dict): op.derive(Cls("MultimodalRoPEOp")) op.derive(Cls("VisionRoPEOp")) op.derive(Cls("QuickGELUOp")) + op.derive(Cls("KimiDeltaAttentionOp")) op.derive(Cls("CopyOp")) op.derive(Cls("CloneOp")) op.derive(Cls("NegOp")) diff --git a/mllm/compile/jit/binary/LinalgIRSerialization.cpp b/mllm/compile/jit/binary/LinalgIRSerialization.cpp index 2a8138853..d005783f7 100644 --- a/mllm/compile/jit/binary/LinalgIRSerialization.cpp +++ b/mllm/compile/jit/binary/LinalgIRSerialization.cpp @@ -39,6 +39,7 @@ #include "mllm/core/aops/STFTOp.hpp" #include "mllm/core/aops/Conv1DOp.hpp" #include "mllm/core/aops/LayerNormOp.hpp" +#include "mllm/core/aops/KimiDeltaAttentionOp.hpp" #include "mllm/compile/jit/binary/LinalgIRSerialization.hpp" namespace mllm::jit::binary { @@ -87,6 +88,7 @@ nlohmann::json dumpLinalgIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { CASE(MultimodalRoPE) CASE(VisionRoPE) CASE(QuickGELU) + CASE(KimiDeltaAttention) CASE(Copy) CASE(Clone) CASE(Neg) @@ -346,6 +348,11 @@ nlohmann::json dumpVisionRoPEOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op nlohmann::json dumpQuickGELUOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { return {}; } +nlohmann::json dumpKimiDeltaAttentionOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { + const auto options = static_cast(op->getAOp())->options(); + return {{"safe_gate", options.safe_gate}, {"lower_bound", options.lower_bound}, {"state_inplace", options.state_inplace}}; +} + nlohmann::json dumpCopyOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { return {}; } nlohmann::json dumpCloneOpIROptions(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 ca3c32d48..e61299d22 100644 --- a/mllm/compile/jit/binary/LinalgIRSerialization.hpp +++ b/mllm/compile/jit/binary/LinalgIRSerialization.hpp @@ -51,6 +51,7 @@ nlohmann::json dumpLayerNormOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) nlohmann::json dumpMultimodalRoPEOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); nlohmann::json dumpVisionRoPEOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); nlohmann::json dumpQuickGELUOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); +nlohmann::json dumpKimiDeltaAttentionOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); nlohmann::json dumpCopyOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); nlohmann::json dumpCloneOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); nlohmann::json dumpNegOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); diff --git a/mllm/compile/jit/interpreter/AopsFromJson.cpp b/mllm/compile/jit/interpreter/AopsFromJson.cpp index f3ecdcbb4..37e655c55 100644 --- a/mllm/compile/jit/interpreter/AopsFromJson.cpp +++ b/mllm/compile/jit/interpreter/AopsFromJson.cpp @@ -32,6 +32,7 @@ #include "mllm/core/aops/LayerNormOp.hpp" #include "mllm/core/aops/VisionRoPEOp.hpp" #include "mllm/core/aops/QuickGELUOp.hpp" +#include "mllm/core/aops/KimiDeltaAttentionOp.hpp" #include "mllm/core/aops/CloneOp.hpp" #include "mllm/core/aops/ConcatOp.hpp" #include "mllm/core/aops/ReduceOps.hpp" @@ -129,6 +130,8 @@ BaseOp::ptr_t aopsFromJson(const nlohmann::json& json) { return __visionRopeFromJson(json); } else if (op_type == "QuickGELU") { return __quickGeluFromJson(json); + } else if (op_type == "KimiDeltaAttention") { + return __kimiDeltaAttentionFromJson(json); } else if (op_type == "Copy") { return __copyFromJson(json); } else if (op_type == "Clone") { @@ -793,6 +796,20 @@ BaseOp::ptr_t __quickGeluFromJson(const nlohmann::json& json) { return op; } +BaseOp::ptr_t __kimiDeltaAttentionFromJson(const nlohmann::json& json) { + aops::KimiDeltaAttentionOpOptions options; + if (json.contains("op_options")) { + const auto& opts = json["op_options"]; + if (opts.contains("safe_gate")) options.safe_gate = opts["safe_gate"]; + if (opts.contains("lower_bound")) options.lower_bound = opts["lower_bound"]; + if (opts.contains("state_inplace")) options.state_inplace = opts["state_inplace"]; + } + + DeviceTypes backend = DeviceTypes::kCPU; + if (json.contains("backend")) { backend = str2DeviceType(json["backend"]); } + return Context::instance().getBackend(backend)->createOp(OpTypes::kKimiDeltaAttention, options); +} + BaseOp::ptr_t __copyFromJson(const nlohmann::json& json) { aops::CopyOpOptions options; diff --git a/mllm/compile/jit/interpreter/AopsFromJson.hpp b/mllm/compile/jit/interpreter/AopsFromJson.hpp index 87d6f2e71..19e07b753 100644 --- a/mllm/compile/jit/interpreter/AopsFromJson.hpp +++ b/mllm/compile/jit/interpreter/AopsFromJson.hpp @@ -47,6 +47,7 @@ BaseOp::ptr_t __layerNormFromJson(const nlohmann::json& json); BaseOp::ptr_t __multimodalRopeFromJson(const nlohmann::json& json); BaseOp::ptr_t __visionRopeFromJson(const nlohmann::json& json); BaseOp::ptr_t __quickGeluFromJson(const nlohmann::json& json); +BaseOp::ptr_t __kimiDeltaAttentionFromJson(const nlohmann::json& json); BaseOp::ptr_t __copyFromJson(const nlohmann::json& json); BaseOp::ptr_t __cloneFromJson(const nlohmann::json& json); BaseOp::ptr_t __negFromJson(const nlohmann::json& json); diff --git a/mllm/core/OpTypes.hpp b/mllm/core/OpTypes.hpp index 76777fe0c..8ec7e95e8 100644 --- a/mllm/core/OpTypes.hpp +++ b/mllm/core/OpTypes.hpp @@ -106,6 +106,7 @@ enum class OpTypes : int32_t { kGroupedQueryAttention = 78, kParallelLinear = 79, kGatedDeltaRule = 80, + kKimiDeltaAttention = 81, // Dynamic Op Start for user to register there own ops. kDynamicOp_Start = 4096, @@ -195,6 +196,7 @@ inline std::string optype2Str(OpTypes type) { case OpTypes::kGroupedQueryAttention: return "GroupedQueryAttention"; case OpTypes::kParallelLinear: return "ParallelLinear"; case OpTypes::kGatedDeltaRule: return "GatedDeltaRule"; + case OpTypes::kKimiDeltaAttention: return "KimiDeltaAttention"; case OpTypes::kDynamicOp_Start: return "DynamicOp_Start"; case OpTypes::kOpType_End: return "OpType_End"; default: return "Unknown"; diff --git a/mllm/core/aops/KimiDeltaAttentionOp.cpp b/mllm/core/aops/KimiDeltaAttentionOp.cpp new file mode 100644 index 000000000..5d065af23 --- /dev/null +++ b/mllm/core/aops/KimiDeltaAttentionOp.cpp @@ -0,0 +1,83 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/core/aops/KimiDeltaAttentionOp.hpp" + +#include +#include + +#include "mllm/compile/ir/linalg/Op.hpp" +#include "mllm/core/Tensor.hpp" +#include "mllm/utils/Common.hpp" + +namespace mllm::aops { + +KimiDeltaAttentionOp::KimiDeltaAttentionOp(const KimiDeltaAttentionOpOptions& options) + : BaseOp(OpTypes::kKimiDeltaAttention), options_(options) {} + +void KimiDeltaAttentionOp::load(const ParameterFile::ptr_t& ploader) { MLLM_EMPTY_SCOPE; } + +void KimiDeltaAttentionOp::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 KimiDeltaAttentionOp::forward(const std::vector& inputs, std::vector& outputs) { + NYI("KimiDeltaAttentionOp::forward not implemented in aops base."); +} + +void KimiDeltaAttentionOp::reshape(const std::vector& inputs, std::vector& outputs) { + if (inputs.size() != 8) { throw std::invalid_argument("KimiDeltaAttention expects 8 input tensors"); } + + const auto& q = inputs[0]; + const auto& k = inputs[1]; + const auto& v = inputs[2]; + const auto& gate = inputs[3]; + const auto& beta = inputs[4]; + const auto& a_log = inputs[5]; + const auto& dt_bias = inputs[6]; + const auto& state = inputs[7]; + if (q.rank() != 4 || k.shape() != q.shape() || v.shape() != q.shape() || gate.shape() != q.shape()) { + throw std::invalid_argument("KimiDeltaAttention q, k, v, and gate must share [B, S, H, D] shape"); + } + const auto batch = q.shape()[0]; + const auto sequence = q.shape()[1]; + const auto heads = q.shape()[2]; + const auto dim = q.shape()[3]; + if (beta.shape() != Tensor::shape_t{batch, sequence, heads}) { + throw std::invalid_argument("KimiDeltaAttention beta must have [B, S, H] shape"); + } + if (a_log.numel() != static_cast(heads) || dt_bias.numel() != static_cast(heads) * dim) { + throw std::invalid_argument("KimiDeltaAttention decay parameters must contain H and H * D values"); + } + if (state.shape() != Tensor::shape_t{batch, heads, dim, dim}) { + throw std::invalid_argument("KimiDeltaAttention state must have [B, H, D, D] shape"); + } + for (const auto& input : inputs) { + if (input.dtype() != kFloat32 || input.device() != q.device()) { + throw std::invalid_argument("KimiDeltaAttention requires float32 inputs on one device"); + } + } + if (options_.safe_gate && (!std::isfinite(options_.lower_bound) || options_.lower_bound >= 0.0F)) { + throw std::invalid_argument("KimiDeltaAttention safe-gate lower bound must be finite and negative"); + } + + outputs.emplace_back(Tensor::empty(q.shape(), q.dtype(), q.device())); + if (options_.state_inplace) { + outputs.emplace_back(state); + } else { + outputs.emplace_back(Tensor::empty(state.shape(), state.dtype(), state.device())); + } +} + +void KimiDeltaAttentionOp::setup(const std::vector& inputs, std::vector& outputs) { + if (options_.state_inplace) { + outputs[0].alloc(); + } else { + BaseOp::setup(inputs, outputs); + } +} + +} // namespace mllm::aops diff --git a/mllm/core/aops/KimiDeltaAttentionOp.hpp b/mllm/core/aops/KimiDeltaAttentionOp.hpp new file mode 100644 index 000000000..44fbe0da4 --- /dev/null +++ b/mllm/core/aops/KimiDeltaAttentionOp.hpp @@ -0,0 +1,33 @@ +// 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 KimiDeltaAttentionOpOptions : public BaseOpOptions { + bool safe_gate = true; + float lower_bound = -5.0F; + bool state_inplace = false; +}; + +class KimiDeltaAttentionOp : public BaseOp { + public: + explicit KimiDeltaAttentionOp(const KimiDeltaAttentionOpOptions& 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 KimiDeltaAttentionOpOptions& options() const { return options_; } + + protected: + KimiDeltaAttentionOpOptions options_; +}; + +} // namespace mllm::aops diff --git a/mllm/models/ling3/configuration_ling3.hpp b/mllm/models/ling3/configuration_ling3.hpp new file mode 100644 index 000000000..65309bb43 --- /dev/null +++ b/mllm/models/ling3/configuration_ling3.hpp @@ -0,0 +1,206 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include "mllm/core/aops/LinearOp.hpp" +#include "mllm/engine/ConfigFile.hpp" + +#include +#include +#include +#include + +namespace mllm::models::ling3 { + +struct Ling3Config : protected ConfigFile { + Ling3Config() = default; + + explicit Ling3Config(const std::string& file_path) : ConfigFile(file_path) { + const auto& config = data(); + hidden_size = config.at("hidden_size"); + intermediate_size = config.at("intermediate_size"); + num_hidden_layers = config.at("num_hidden_layers"); + num_attention_heads = config.at("num_attention_heads"); + num_key_value_heads = config.at("num_key_value_heads"); + head_dim = config.at("head_dim"); + vocab_size = config.at("vocab_size"); + max_position_embeddings = config.at("max_position_embeddings"); + rms_norm_eps = config.at("rms_norm_eps"); + rope_theta = config.at("rope_theta"); + hidden_act = config.at("hidden_act"); + tie_word_embeddings = config.at("tie_word_embeddings"); + use_qkv_bias = config.at("use_qkv_bias"); + + layer_group_size = config.at("layer_group_size"); + short_conv_kernel_size = config.at("short_conv_kernel_size"); + no_kda_lora = config.at("no_kda_lora"); + kda_safe_gate = config.at("kda_safe_gate"); + kda_lower_bound = config.at("kda_lower_bound"); + + q_lora_rank = config.at("q_lora_rank"); + kv_lora_rank = config.at("kv_lora_rank"); + qk_rope_head_dim = config.at("qk_rope_head_dim"); + qk_nope_head_dim = config.at("qk_nope_head_dim"); + qk_head_dim = config.at("qk_head_dim"); + v_head_dim = config.at("v_head_dim"); + rope_interleave = config.at("rope_interleave"); + gated_attention_proj_granularity_type = config.at("gated_attention_proj_granularity_type"); + + num_experts = config.at("num_experts"); + num_shared_experts = config.at("num_shared_experts"); + num_experts_per_tok = config.at("num_experts_per_tok"); + n_group = config.at("n_group"); + topk_group = config.at("topk_group"); + moe_intermediate_size = config.at("moe_intermediate_size"); + moe_shared_expert_intermediate_size = config.at("moe_shared_expert_intermediate_size"); + first_k_dense_replace = config.at("first_k_dense_replace"); + routed_scaling_factor = config.at("routed_scaling_factor"); + scoring_func = config.at("scoring_func"); + topk_method = config.at("topk_method"); + moe_router_enable_expert_bias = config.at("moe_router_enable_expert_bias"); + + pad_token_id = config.at("pad_token_id"); + eos_token_id = config.at("eos_token_id"); + max_cache_length = config.value("max_cache_length", max_cache_length); + if (config.contains("linear_impl_type")) { + const auto impl_name = config.at("linear_impl_type").get(); + linear_impl_type = aops::str2LinearImplTypes(impl_name); + if (linear_impl_type == aops::LinearImplTypes::kDefault && impl_name != "Default") { + throw std::invalid_argument("Ling-3 contains an unsupported linear_impl_type: " + impl_name); + } + } + + validate(); + } + + int32_t hidden_size = 1536; + int32_t intermediate_size = 4608; + int32_t num_hidden_layers = 24; + int32_t num_attention_heads = 16; + int32_t num_key_value_heads = 16; + int32_t head_dim = 128; + int32_t vocab_size = 157184; + int32_t max_position_embeddings = 131072; + float rms_norm_eps = 1.0e-6F; + float rope_theta = 6000000.0F; + std::string hidden_act = "silu"; + bool tie_word_embeddings = false; + bool use_qkv_bias = false; + + int32_t layer_group_size = 4; + int32_t short_conv_kernel_size = 4; + bool no_kda_lora = true; + bool kda_safe_gate = true; + float kda_lower_bound = -5.0F; + + int32_t q_lora_rank = 256; + int32_t kv_lora_rank = 512; + int32_t qk_rope_head_dim = 64; + int32_t qk_nope_head_dim = 128; + int32_t qk_head_dim = 192; + int32_t v_head_dim = 128; + bool rope_interleave = true; + std::string gated_attention_proj_granularity_type = "head_wise"; + + int32_t num_experts = 128; + int32_t num_shared_experts = 1; + int32_t num_experts_per_tok = 8; + int32_t n_group = 8; + int32_t topk_group = 4; + int32_t moe_intermediate_size = 512; + int32_t moe_shared_expert_intermediate_size = 512; + int32_t first_k_dense_replace = 1; + float routed_scaling_factor = 2.5F; + std::string scoring_func = "sigmoid"; + std::string topk_method = "noaux_tc"; + bool moe_router_enable_expert_bias = true; + + int64_t pad_token_id = 156892; + int64_t eos_token_id = 156895; + int64_t end_of_text_token_id = 156892; + int64_t bos_token_id = 156891; + int32_t max_cache_length = 2048; + + aops::LinearImplTypes linear_impl_type = aops::LinearImplTypes::kDefault; + + [[nodiscard]] bool isFullAttentionLayer(int32_t layer_index) const { + return (layer_index + 1) % layer_group_size == 0 || layer_index >= num_hidden_layers / layer_group_size * layer_group_size; + } + + [[nodiscard]] int32_t numFullAttentionLayers() const { + int32_t count = 0; + for (int32_t layer = 0; layer < num_hidden_layers; ++layer) { + if (isFullAttentionLayer(layer)) { ++count; } + } + return count; + } + + [[nodiscard]] int32_t numKDALayers() const { return num_hidden_layers - numFullAttentionLayers(); } + + 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 || vocab_size <= 0 || max_position_embeddings <= 0 + || max_cache_length <= 0) { + throw std::invalid_argument("Ling-3 contains invalid model dimensions"); + } + if (layer_group_size <= 0 || short_conv_kernel_size <= 1 || q_lora_rank <= 0 || kv_lora_rank <= 0 || qk_rope_head_dim <= 0 + || qk_nope_head_dim <= 0 || qk_head_dim != qk_nope_head_dim + qk_rope_head_dim || v_head_dim <= 0) { + throw std::invalid_argument("Ling-3 contains invalid hybrid attention dimensions"); + } + if (num_experts <= 0 || num_shared_experts != 1 || num_experts_per_tok <= 0 || num_experts_per_tok > num_experts + || n_group <= 0 || num_experts % n_group != 0 || topk_group <= 0 || topk_group > n_group || moe_intermediate_size <= 0 + || moe_shared_expert_intermediate_size <= 0 || first_k_dense_replace < 0 || first_k_dense_replace > num_hidden_layers) { + throw std::invalid_argument("Ling-3 contains invalid MoE dimensions"); + } + if (!std::isfinite(rms_norm_eps) || rms_norm_eps <= 0.0F || !std::isfinite(rope_theta) || rope_theta <= 0.0F + || !std::isfinite(kda_lower_bound) || kda_lower_bound >= 0.0F || !std::isfinite(routed_scaling_factor) + || routed_scaling_factor <= 0.0F) { + throw std::invalid_argument("Ling-3 contains invalid numerical parameters"); + } + if (hidden_act != "silu" || scoring_func != "sigmoid" || topk_method != "noaux_tc" || !no_kda_lora || !kda_safe_gate + || !rope_interleave || gated_attention_proj_granularity_type != "head_wise" || tie_word_embeddings || use_qkv_bias + || !moe_router_enable_expert_bias) { + throw std::invalid_argument("Ling-3 CPU received an unsupported architecture variant"); + } + } +}; + +inline bool hasOfficialLing3TinyArchitecture(const Ling3Config& config) { + return config.hidden_size == 1536 && config.intermediate_size == 4608 && config.num_hidden_layers == 24 + && config.num_attention_heads == 16 && config.num_key_value_heads == 16 && config.head_dim == 128 + && config.vocab_size == 157184 && config.max_position_embeddings == 131072 && config.rms_norm_eps == 1.0e-6F + && config.rope_theta == 6000000.0F && config.layer_group_size == 4 && config.numFullAttentionLayers() == 6 + && config.numKDALayers() == 18 && config.short_conv_kernel_size == 4 && config.q_lora_rank == 256 + && config.kv_lora_rank == 512 && config.qk_rope_head_dim == 64 && config.qk_nope_head_dim == 128 + && config.qk_head_dim == 192 && config.v_head_dim == 128 && config.num_experts == 128 + && config.num_experts_per_tok == 8 && config.n_group == 8 && config.topk_group == 4 + && config.moe_intermediate_size == 512 && config.moe_shared_expert_intermediate_size == 512 + && config.first_k_dense_replace == 1 && config.routed_scaling_factor == 2.5F && config.pad_token_id == 156892 + && config.eos_token_id == 156895; +} + +inline void validateLing3ModelConfigMatch(const Ling3Config& config, const ParameterFile::ptr_t& parameter_file) { + constexpr auto kEmbeddingWeight = "model.word_embeddings.weight"; + constexpr auto kLmHeadWeight = "lm_head.weight"; + if (!hasOfficialLing3TinyArchitecture(config)) { + throw std::invalid_argument("Ling-3 model/config mismatch: expected inclusionAI/Ling-3.0-tiny"); + } + if (parameter_file == nullptr || parameter_file->version() != ModelFileVersion::kV2) { + throw std::invalid_argument("Ling-3 mobile CPU inference requires an MLLM V2 model file"); + } + if (!parameter_file->has(kEmbeddingWeight) || !parameter_file->has(kLmHeadWeight)) { + throw std::invalid_argument("Ling-3 model file is missing untied embedding or LM-head weights"); + } + const auto embedding = parameter_file->pull(kEmbeddingWeight); + if (embedding.dtype() != kFloat32 || embedding.shape() != Tensor::shape_t({config.vocab_size, config.hidden_size})) { + throw std::invalid_argument("Ling-3 model file has an incompatible word embedding descriptor"); + } + const auto lm_head = parameter_file->pull(kLmHeadWeight); + if (lm_head.dtype() != kByte || lm_head.shape().size() != 1) { + throw std::invalid_argument("Ling-3 mobile model requires a packed KAI W4A32 LM head"); + } +} + +} // namespace mllm::models::ling3 diff --git a/mllm/models/ling3/modeling_ling3.hpp b/mllm/models/ling3/modeling_ling3.hpp new file mode 100644 index 000000000..9c478bb5a --- /dev/null +++ b/mllm/models/ling3/modeling_ling3.hpp @@ -0,0 +1,663 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include "mllm/mllm.hpp" +#include "mllm/models/ARGeneration.hpp" +#include "mllm/models/ling3/configuration_ling3.hpp" +#include "mllm/nn/Functional.hpp" +#include "mllm/nn/Module.hpp" +#include "mllm/nn/Nn.hpp" +#include "mllm/nn/lmcache/StaticCache.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mllm::models::ling3 { + +// Model-local tensor preparation: Ling-3 uses an adjacent-pair checkpoint RoPE +// layout that must be converted to mllm's half-split attention layout. These +// helpers prepare RoPE/cache tensors only; reusable compute is dispatched +// through registered nn layers and functional operations below. +inline Tensor makeLing3RoPEInvFreq(int rotary_dim, float rope_theta) { + auto inv_freq = Tensor::empty({rotary_dim / 2}, kFloat32, kCPU).alloc(); + auto* values = inv_freq.ptr(); + for (int index = 0; index < rotary_dim / 2; ++index) { + values[index] = 1.0F / std::pow(rope_theta, 2.0F * static_cast(index) / static_cast(rotary_dim)); + } + return inv_freq; +} + +inline std::pair makeLing3RotaryEmbedding(const Tensor& position_ids, const Tensor& inv_freq) { + if (position_ids.shape().size() != 2 || position_ids.dtype() != kInt64 || inv_freq.shape().size() != 1 + || inv_freq.dtype() != kFloat32) { + throw std::invalid_argument("Ling-3 RoPE received invalid position ids or frequencies"); + } + const int batch = position_ids.shape()[0]; + const int sequence = position_ids.shape()[1]; + const int half_dim = inv_freq.shape()[0]; + const int rotary_dim = half_dim * 2; + auto cos = Tensor::empty({batch, sequence, rotary_dim}, kFloat32, kCPU).alloc(); + auto sin = Tensor::empty({batch, sequence, rotary_dim}, kFloat32, kCPU).alloc(); + const auto* positions = position_ids.ptr(); + const auto* frequencies = inv_freq.ptr(); + auto* cos_values = cos.ptr(); + auto* sin_values = sin.ptr(); + for (int b = 0; b < batch; ++b) { + for (int s = 0; s < sequence; ++s) { + for (int dim = 0; dim < half_dim; ++dim) { + const float angle = static_cast(positions[b * sequence + s]) * frequencies[dim]; + const float cos_value = std::cos(angle); + const float sin_value = std::sin(angle); + const std::size_t base = (static_cast(b) * sequence + s) * rotary_dim; + cos_values[base + dim] = cos_value; + cos_values[base + half_dim + dim] = cos_value; + sin_values[base + dim] = sin_value; + sin_values[base + half_dim + dim] = sin_value; + } + } + } + return {cos, sin}; +} + +// Converts adjacent-pair RoPE storage [x0, x1, x2, x3, ...] to the +// half-split layout used by the official Bailing implementation while applying +// the rotation. Input/output are [B, H, S, D]; cos/sin are [B, S, D]. +inline Tensor applyLing3InterleavedRoPE(Tensor input, const Tensor& cos, const Tensor& sin) { + if (input.dtype() != kFloat32 || cos.dtype() != kFloat32 || sin.dtype() != kFloat32 || input.shape().size() != 4 + || cos.shape().size() != 3 || sin.shape() != cos.shape()) { + throw std::invalid_argument("Ling-3 interleaved RoPE currently requires float32 tensors"); + } + const int batch = input.shape()[0]; + const int heads = input.shape()[1]; + const int sequence = input.shape()[2]; + const int dim = input.shape()[3]; + if (dim <= 0 || dim % 2 != 0 || cos.shape()[0] != batch || cos.shape()[1] != sequence || cos.shape()[2] != dim) { + throw std::invalid_argument("Ling-3 interleaved RoPE received incompatible shapes"); + } + auto source = input.contiguous(); + auto output = Tensor::empty(input.shape(), kFloat32, kCPU).alloc(); + const auto* source_values = source.ptr(); + const auto* cos_values = cos.ptr(); + const auto* sin_values = sin.ptr(); + auto* output_values = output.ptr(); + const int half_dim = dim / 2; + for (int b = 0; b < batch; ++b) { + for (int h = 0; h < heads; ++h) { + for (int s = 0; s < sequence; ++s) { + const std::size_t vector_base = ((static_cast(b) * heads + h) * sequence + s) * dim; + const std::size_t rope_base = (static_cast(b) * sequence + s) * dim; + for (int pair = 0; pair < half_dim; ++pair) { + const float even = source_values[vector_base + pair * 2]; + const float odd = source_values[vector_base + pair * 2 + 1]; + const float cos_value = cos_values[rope_base + pair]; + const float sin_value = sin_values[rope_base + pair]; + output_values[vector_base + pair] = even * cos_value - odd * sin_value; + output_values[vector_base + half_dim + pair] = odd * cos_value + even * sin_value; + } + } + } + } + return output; +} + +inline Tensor padLing3ValuesForCache(Tensor input, int output_dim) { + if (input.dtype() != kFloat32 || input.shape().size() != 4 || input.shape()[3] <= 0 || input.shape()[3] > output_dim) { + throw std::invalid_argument("Ling-3 MLA value-cache padding received an invalid tensor"); + } + auto source = input.contiguous(); + auto output = Tensor::zeros({input.shape()[0], input.shape()[1], input.shape()[2], output_dim}, kFloat32, kCPU); + const int input_dim = input.shape()[3]; + const std::size_t vectors = static_cast(input.shape()[0]) * input.shape()[1] * input.shape()[2]; + const auto* source_values = source.ptr(); + auto* output_values = output.ptr(); + for (std::size_t vector = 0; vector < vectors; ++vector) { + std::memcpy(output_values + vector * output_dim, source_values + vector * input_dim, + static_cast(input_dim) * sizeof(float)); + } + return output; +} + +class Ling3MLP final : public nn::Module { + public: + Ling3MLP() = default; + Ling3MLP(const std::string& name, const Ling3Config& config, const std::optional& intermediate_size = std::nullopt) + : nn::Module(name) { + const int32_t intermediate = intermediate_size.value_or(config.intermediate_size); + gate_proj_ = reg("gate_proj", config.hidden_size, intermediate, false, config.linear_impl_type); + up_proj_ = reg("up_proj", config.hidden_size, intermediate, false, config.linear_impl_type); + down_proj_ = reg("down_proj", intermediate, config.hidden_size, false, config.linear_impl_type); + activation_ = reg("act"); + } + + std::vector forward(const std::vector& inputs, const std::vector& args) override { + return {down_proj_(activation_(gate_proj_(inputs[0])) * up_proj_(inputs[0]))}; + } + + private: + nn::Linear gate_proj_; + nn::Linear up_proj_; + nn::Linear down_proj_; + nn::SiLU activation_; +}; + +class Ling3MoEGate final : public nn::Module { + public: + Ling3MoEGate() = default; + Ling3MoEGate(const std::string& name, const Ling3Config& config) : nn::Module(name) { + hidden_size_ = config.hidden_size; + num_experts_ = config.num_experts; + top_k_ = config.num_experts_per_tok; + num_groups_ = config.n_group; + top_groups_ = config.topk_group; + routed_scaling_factor_ = config.routed_scaling_factor; + weight_ = reg("weight", getModuleName() + ".weight"); + expert_bias_ = reg("expert_bias", getModuleName() + ".expert_bias"); + } + + std::vector forward(const std::vector& inputs, const std::vector& args) override { + auto hidden = inputs[0]; + if (hidden.dtype() != kFloat32 || hidden.shape().size() != 3 || hidden.shape()[2] != hidden_size_ + || weight_.weight().dtype() != kFloat32 || expert_bias_.weight().dtype() != kFloat32) { + throw std::invalid_argument("Ling-3 MoE router requires float32 tensors with official shapes"); + } + // Model-level orchestration: the official Ling grouped top-k rule includes + // expert bias and a top-2-per-group score, which the existing generic mllm + // MoE layers do not represent. Matmul and every expert MLP remain normal + // nn operations; this loop only produces routing ids and weights. + hidden = hidden.view({-1, hidden_size_}); + auto logits = nn::functional::matmul(hidden, weight_.weight(), false, true).contiguous(); + const int tokens = hidden.shape()[0]; + auto topk_ids = Tensor::empty({tokens, top_k_}, kInt32, kCPU).alloc(); + auto topk_weights = Tensor::empty({tokens, top_k_}, kFloat32, kCPU).alloc(); + const auto* logits_values = logits.ptr(); + const auto* bias_values = expert_bias_.weight().ptr(); + auto* id_values = topk_ids.ptr(); + auto* weight_values = topk_weights.ptr(); + const int experts_per_group = num_experts_ / num_groups_; + + std::vector scores(static_cast(num_experts_)); + std::vector> ranked_groups(static_cast(num_groups_)); + std::vector> ranked_experts; + ranked_experts.reserve(static_cast(top_groups_ * experts_per_group)); + for (int token = 0; token < tokens; ++token) { + for (int expert = 0; expert < num_experts_; ++expert) { + const float value = logits_values[static_cast(token) * num_experts_ + expert]; + scores[expert] = value >= 0.0F ? 1.0F / (1.0F + std::exp(-value)) : std::exp(value) / (1.0F + std::exp(value)); + } + for (int group = 0; group < num_groups_; ++group) { + float largest = -std::numeric_limits::infinity(); + float second_largest = -std::numeric_limits::infinity(); + for (int offset = 0; offset < experts_per_group; ++offset) { + const int expert = group * experts_per_group + offset; + const float routed_score = scores[expert] + bias_values[expert]; + if (routed_score > largest) { + second_largest = largest; + largest = routed_score; + } else if (routed_score > second_largest) { + second_largest = routed_score; + } + } + ranked_groups[group] = {largest + second_largest, group}; + } + std::partial_sort(ranked_groups.begin(), ranked_groups.begin() + top_groups_, ranked_groups.end(), + [](const auto& lhs, const auto& rhs) { + return lhs.first != rhs.first ? lhs.first > rhs.first : lhs.second < rhs.second; + }); + ranked_experts.clear(); + for (int group_index = 0; group_index < top_groups_; ++group_index) { + const int group = ranked_groups[group_index].second; + for (int offset = 0; offset < experts_per_group; ++offset) { + const int expert = group * experts_per_group + offset; + ranked_experts.emplace_back(scores[expert] + bias_values[expert], expert); + } + } + std::partial_sort(ranked_experts.begin(), ranked_experts.begin() + top_k_, ranked_experts.end(), + [](const auto& lhs, const auto& rhs) { + return lhs.first != rhs.first ? lhs.first > rhs.first : lhs.second < rhs.second; + }); + float score_sum = 1.0e-20F; + for (int route = 0; route < top_k_; ++route) { score_sum += scores[ranked_experts[route].second]; } + for (int route = 0; route < top_k_; ++route) { + const int expert = ranked_experts[route].second; + id_values[token * top_k_ + route] = expert; + weight_values[token * top_k_ + route] = scores[expert] / score_sum * routed_scaling_factor_; + } + } + return {topk_ids, topk_weights}; + } + + private: + int32_t hidden_size_ = 0; + int32_t num_experts_ = 0; + int32_t top_k_ = 0; + int32_t num_groups_ = 0; + int32_t top_groups_ = 0; + float routed_scaling_factor_ = 1.0F; + nn::Param weight_; + nn::Param expert_bias_; +}; + +class Ling3SparseMoE final : public nn::Module { + public: + Ling3SparseMoE() = default; + Ling3SparseMoE(const std::string& name, const Ling3Config& config) : nn::Module(name) { + top_k_ = config.num_experts_per_tok; + experts_ = reg>("experts", config.num_experts, config, + std::optional(config.moe_intermediate_size)); + gate_ = reg("gate", config); + shared_experts_ = + reg("shared_experts", config, + std::optional(config.moe_shared_expert_intermediate_size * config.num_shared_experts)); + } + + std::vector forward(const std::vector& inputs, const std::vector& args) override { + auto identity = inputs[0]; + const auto original_shape = identity.shape(); + auto routing = gate_(identity); + auto hidden = identity.view({-1, identity.shape().back()}); + auto routed = moeInfer(hidden, routing[0], routing[1]).view(original_shape); + return {routed + shared_experts_(identity)[0]}; + } + + private: + Tensor moeInfer(const Tensor& input, Tensor& topk_ids, Tensor& topk_weights) { + const int tokens = topk_ids.shape()[0]; + const int hidden_size = input.shape()[1]; + auto routed_output = Tensor::empty({tokens, top_k_, hidden_size}, kFloat32, kCPU).alloc(); + const auto* id_values = topk_ids.ptr(); + auto* routed_values = routed_output.ptr(); + // Model-level sparse dispatch: keep each routed expert invocation at M=1. + // KAI's dynamic activation + // quantization is row-local, and this avoids platform-dependent grouped + // prefill results while retaining the same token/expert routing contract. + for (int token = 0; token < tokens; ++token) { + auto expert_input = input[{{token, token + 1}, kAll}]; + for (int route = 0; route < top_k_; ++route) { + const int expert = id_values[token * top_k_ + route]; + auto expert_output = experts_.list()[expert](expert_input)[0].contiguous(); + std::memcpy(routed_values + (static_cast(token) * top_k_ + route) * hidden_size, + expert_output.ptr(), static_cast(hidden_size) * sizeof(float)); + } + } + return routed_output.mul_(topk_weights.unsqueeze(-1)).sum(1).to(routed_output.dtype()); + } + + int32_t top_k_ = 0; + nn::ModuleList experts_; + Ling3MoEGate gate_; + Ling3MLP shared_experts_; +}; + +class Ling3KimiDeltaAttention final : public nn::Module { + public: + // The official checkpoint was validated against the current-first + // accumulation order (current tap first, then the K - 1 history taps). + static constexpr auto kConvAccumulationOrder = aops::CausalDepthwiseConv1DAccumulationOrder::kCurrentFirst; + + Ling3KimiDeltaAttention() = default; + Ling3KimiDeltaAttention(const std::string& name, const Ling3Config& config) : nn::Module(name) { + hidden_size_ = config.hidden_size; + num_heads_ = config.num_attention_heads; + head_dim_ = config.head_dim; + projection_size_ = num_heads_ * head_dim_; + conv_size_ = config.short_conv_kernel_size; + safe_gate_ = config.kda_safe_gate; + lower_bound_ = config.kda_lower_bound; + + q_proj_ = reg("q_proj", hidden_size_, projection_size_, false, config.linear_impl_type); + k_proj_ = reg("k_proj", hidden_size_, projection_size_, false, config.linear_impl_type); + v_proj_ = reg("v_proj", hidden_size_, projection_size_, false, config.linear_impl_type); + // Stateful depthwise causal short convolutions. The registered operation + // owns the official `{q,k,v}_conv1d.weight` parameters and the explicit + // [B, C, K - 1] history; the current-first accumulation order reuses the + // existing GDN causal-convolution kernel. + q_conv1d_ = reg("q_conv1d", projection_size_, conv_size_, /*bias=*/false, + /*state_inplace=*/true, kConvAccumulationOrder); + k_conv1d_ = reg("k_conv1d", projection_size_, conv_size_, /*bias=*/false, + /*state_inplace=*/true, kConvAccumulationOrder); + v_conv1d_ = reg("v_conv1d", projection_size_, conv_size_, /*bias=*/false, + /*state_inplace=*/true, kConvAccumulationOrder); + f_proj_ = reg("f_proj", hidden_size_, projection_size_, false, config.linear_impl_type); + b_proj_ = reg("b_proj", hidden_size_, num_heads_, false, config.linear_impl_type); + g_proj_ = reg("g_proj", hidden_size_, projection_size_, false, config.linear_impl_type); + A_log_ = reg("A_log", getModuleName() + ".A_log"); + dt_bias_ = reg("dt_bias", getModuleName() + ".dt_bias"); + o_norm_ = reg("o_norm", config.rms_norm_eps, false); + o_proj_ = reg("o_proj", projection_size_, hidden_size_, false, config.linear_impl_type); + silu_ = reg("conv_activation"); + sigmoid_ = reg("gate_activation"); + kda_ = reg("kda", safe_gate_, lower_bound_, /*state_inplace=*/true); + } + + void resetState(int batch_size) { + if (batch_size <= 0) { throw std::invalid_argument("Ling-3 KDA reset requires a positive batch size"); } + // Module-owned request state. State transitions themselves are explicit + // outputs of the registered causal-convolution and KDA operations. + recurrent_state_ = Tensor::zeros({batch_size, num_heads_, head_dim_, head_dim_}, kFloat32, kCPU); + q_conv_state_ = Tensor::zeros({batch_size, projection_size_, conv_size_ - 1}, kFloat32, kCPU); + k_conv_state_ = Tensor::zeros({batch_size, projection_size_, conv_size_ - 1}, kFloat32, kCPU); + v_conv_state_ = Tensor::zeros({batch_size, projection_size_, conv_size_ - 1}, kFloat32, kCPU); + } + + std::vector forward(const std::vector& inputs, const std::vector& args) override { + auto hidden = inputs[0]; + const int batch = hidden.shape()[0]; + const int sequence = hidden.shape()[1]; + if (recurrent_state_.isNil() || recurrent_state_.shape()[0] != batch) { resetState(batch); } + + auto q = q_proj_(hidden).contiguous(); + auto k = k_proj_(hidden).contiguous(); + auto v = v_proj_(hidden).contiguous(); + auto gate_logits = f_proj_(hidden).contiguous(); + auto beta = sigmoid_(b_proj_(hidden)).contiguous(); + const auto a_log = A_log_.weight(); + const auto dt_bias = dt_bias_.weight(); + if (q.dtype() != kFloat32 || k.dtype() != kFloat32 || v.dtype() != kFloat32 || gate_logits.dtype() != kFloat32 + || beta.dtype() != kFloat32 || a_log.dtype() != kFloat32 || dt_bias.dtype() != kFloat32) { + throw std::invalid_argument("Ling-3 KDA requires float32 recurrent activations and parameters"); + } + + auto [q_conv, updated_q_conv_state] = q_conv1d_(q, q_conv_state_); + auto [k_conv, updated_k_conv_state] = k_conv1d_(k, k_conv_state_); + auto [v_conv, updated_v_conv_state] = v_conv1d_(v, v_conv_state_); + q_conv_state_ = std::move(updated_q_conv_state); + k_conv_state_ = std::move(updated_k_conv_state); + v_conv_state_ = std::move(updated_v_conv_state); + q = silu_(q_conv).view({batch, sequence, num_heads_, head_dim_}).contiguous(); + k = silu_(k_conv).view({batch, sequence, num_heads_, head_dim_}).contiguous(); + v = silu_(v_conv).view({batch, sequence, num_heads_, head_dim_}).contiguous(); + gate_logits = gate_logits.view({batch, sequence, num_heads_, head_dim_}).contiguous(); + auto [output, updated_state] = kda_(q, k, v, gate_logits, beta, a_log, dt_bias, recurrent_state_); + recurrent_state_ = std::move(updated_state); + + output = output.view({batch * sequence * num_heads_, head_dim_}); + auto output_gate = sigmoid_(g_proj_(hidden).view({batch * sequence * num_heads_, head_dim_})); + output = o_norm_(output) * output_gate; + output = output.view({batch, sequence, projection_size_}); + return {o_proj_(output)}; + } + + private: + int32_t hidden_size_ = 0; + int32_t num_heads_ = 0; + int32_t head_dim_ = 0; + int32_t projection_size_ = 0; + int32_t conv_size_ = 0; + bool safe_gate_ = true; + float lower_bound_ = -5.0F; + nn::Linear q_proj_, k_proj_, v_proj_; + nn::CausalDepthwiseConv1D q_conv1d_, k_conv1d_, v_conv1d_; + nn::Linear f_proj_, b_proj_, g_proj_; + nn::Param A_log_, dt_bias_; + nn::RMSNorm o_norm_; + nn::Linear o_proj_; + nn::SiLU silu_; + nn::Sigmoid sigmoid_; + nn::KimiDeltaAttention kda_; + Tensor recurrent_state_, q_conv_state_, k_conv_state_, v_conv_state_; +}; + +class Ling3MultiLatentAttention final : public nn::Module { + public: + int32_t cache_layer_index_ = 0; + + Ling3MultiLatentAttention() = default; + Ling3MultiLatentAttention(const std::string& name, const Ling3Config& config) : nn::Module(name) { + hidden_size_ = config.hidden_size; + num_heads_ = config.num_attention_heads; + q_lora_rank_ = config.q_lora_rank; + kv_lora_rank_ = config.kv_lora_rank; + qk_rope_dim_ = config.qk_rope_head_dim; + qk_nope_dim_ = config.qk_nope_head_dim; + qk_dim_ = config.qk_head_dim; + value_dim_ = config.v_head_dim; + q_a_proj_ = reg("q_a_proj", hidden_size_, q_lora_rank_, config.use_qkv_bias, config.linear_impl_type); + q_a_layernorm_ = reg("q_a_layernorm", config.rms_norm_eps, false); + q_b_proj_ = reg("q_b_proj", q_lora_rank_, num_heads_ * qk_dim_, false, config.linear_impl_type); + kv_a_proj_ = reg("kv_a_proj_with_mqa", hidden_size_, kv_lora_rank_ + qk_rope_dim_, config.use_qkv_bias, + config.linear_impl_type); + kv_a_layernorm_ = reg("kv_a_layernorm", config.rms_norm_eps, false); + kv_b_proj_ = + reg("kv_b_proj", kv_lora_rank_, num_heads_ * (qk_nope_dim_ + value_dim_), false, config.linear_impl_type); + g_proj_ = reg("g_proj", hidden_size_, num_heads_, false, config.linear_impl_type); + dense_ = reg("dense", num_heads_ * value_dim_, hidden_size_, config.use_qkv_bias, config.linear_impl_type); + mask_ = reg("mask"); + softmax_ = reg("softmax", -1); + sigmoid_ = reg("gate_sigmoid"); + } + + std::vector forward(const std::vector& inputs, const std::vector& args) override { + auto hidden = inputs[0]; + const auto cos = inputs[1]; + const auto sin = inputs[2]; + auto* cache = args[0].get(); + const int batch = hidden.shape()[0]; + const int sequence = hidden.shape()[1]; + + auto query = q_b_proj_(q_a_layernorm_(q_a_proj_(hidden))).view({batch, sequence, num_heads_, qk_dim_}).transpose(1, 2); + auto q_nope = query[{kAll, kAll, kAll, {0, qk_nope_dim_}}].contiguous(); + auto q_rope = query[{kAll, kAll, kAll, {qk_nope_dim_, qk_dim_}}].contiguous(); + + auto compressed_kv = kv_a_proj_(hidden); + auto kv_latent = compressed_kv[{kAll, kAll, {0, kv_lora_rank_}}].contiguous(); + auto k_rope = compressed_kv[{kAll, kAll, {kv_lora_rank_, kv_lora_rank_ + qk_rope_dim_}}] + .contiguous() + .view({batch, sequence, 1, qk_rope_dim_}) + .transpose(1, 2); + auto kv_expanded = + kv_b_proj_(kv_a_layernorm_(kv_latent)).view({batch, sequence, num_heads_, qk_nope_dim_ + value_dim_}).transpose(1, 2); + auto k_nope = kv_expanded[{kAll, kAll, kAll, {0, qk_nope_dim_}}].contiguous(); + auto value = kv_expanded[{kAll, kAll, kAll, {qk_nope_dim_, qk_nope_dim_ + value_dim_}}].contiguous(); + + q_rope = applyLing3InterleavedRoPE(q_rope, cos, sin); + k_rope = applyLing3InterleavedRoPE(k_rope, cos, sin).repeat(num_heads_, 1); + query = nn::functional::concat({q_nope, q_rope}, -1); + auto key = nn::functional::concat({k_nope, k_rope}, -1); + auto padded_value = padLing3ValuesForCache(value, qk_dim_); + auto cached = cache->updateKVCache(cache_layer_index_, key, padded_value); + key = cached[0]; + padded_value = cached[1]; + + auto attention = nn::functional::matmul(query, key, false, true) * (1.0F / std::sqrt(static_cast(qk_dim_))); + attention = softmax_(mask_(attention)); + auto output = nn::functional::matmul(attention, padded_value); + output = output[{kAll, kAll, kAll, {0, value_dim_}}].contiguous().transpose(1, 2); + auto gate = sigmoid_(g_proj_(hidden)).unsqueeze(-1); + output = output * gate; + output = output.view({batch, sequence, num_heads_ * value_dim_}); + return {dense_(output)}; + } + + private: + int32_t hidden_size_ = 0, num_heads_ = 0, q_lora_rank_ = 0, kv_lora_rank_ = 0; + int32_t qk_rope_dim_ = 0, qk_nope_dim_ = 0, qk_dim_ = 0, value_dim_ = 0; + nn::Linear q_a_proj_, q_b_proj_, kv_a_proj_, kv_b_proj_, g_proj_, dense_; + nn::RMSNorm q_a_layernorm_, kv_a_layernorm_; + nn::CausalMask mask_; + nn::Softmax softmax_; + nn::Sigmoid sigmoid_; +}; + +class Ling3KDAdecoder final : public nn::Module { + public: + Ling3KimiDeltaAttention attention_; + + Ling3KDAdecoder() = default; + Ling3KDAdecoder(const std::string& name, const Ling3Config& config, int32_t layer_index) : nn::Module(name) { + attention_ = reg("attention", config); + input_layernorm_ = reg("input_layernorm", config.rms_norm_eps, false); + post_attention_layernorm_ = reg("post_attention_layernorm", config.rms_norm_eps, false); + use_moe_ = layer_index >= config.first_k_dense_replace; + if (use_moe_) { + moe_ = reg("mlp", config); + } else { + dense_mlp_ = reg("mlp", config); + } + } + + std::vector forward(const std::vector& inputs, const std::vector& args) override { + auto hidden = inputs[0]; + auto attention_output = attention_(input_layernorm_(hidden))[0]; + hidden = hidden + attention_output; + auto mlp_input = post_attention_layernorm_(hidden); + return {hidden + (use_moe_ ? moe_(mlp_input)[0] : dense_mlp_(mlp_input)[0])}; + } + + private: + bool use_moe_ = false; + nn::RMSNorm input_layernorm_, post_attention_layernorm_; + Ling3MLP dense_mlp_; + Ling3SparseMoE moe_; +}; + +class Ling3MLADecoder final : public nn::Module { + public: + Ling3MultiLatentAttention attention_; + + Ling3MLADecoder() = default; + Ling3MLADecoder(const std::string& name, const Ling3Config& config, int32_t layer_index) : nn::Module(name) { + attention_ = reg("attention", config); + input_layernorm_ = reg("input_layernorm", config.rms_norm_eps, false); + post_attention_layernorm_ = reg("post_attention_layernorm", config.rms_norm_eps, false); + use_moe_ = layer_index >= config.first_k_dense_replace; + if (use_moe_) { + moe_ = reg("mlp", config); + } else { + dense_mlp_ = reg("mlp", config); + } + } + + std::vector forward(const std::vector& inputs, const std::vector& args) override { + auto hidden = inputs[0]; + auto attention_output = attention_(input_layernorm_(hidden), inputs[1], inputs[2], args[0])[0]; + hidden = hidden + attention_output; + auto mlp_input = post_attention_layernorm_(hidden); + return {hidden + (use_moe_ ? moe_(mlp_input)[0] : dense_mlp_(mlp_input)[0])}; + } + + private: + bool use_moe_ = false; + nn::RMSNorm input_layernorm_, post_attention_layernorm_; + Ling3MLP dense_mlp_; + Ling3SparseMoE moe_; +}; + +class Ling3Model final : public nn::Module { + public: + Ling3Model() = default; + Ling3Model(const std::string& name, const Ling3Config& config) : nn::Module(name) { + word_embeddings_ = reg("word_embeddings", config.vocab_size, config.hidden_size); + int32_t kda_index = 0; + int32_t mla_index = 0; + for (int32_t layer = 0; layer < config.num_hidden_layers; ++layer) { + const std::string layer_name = "layers." + std::to_string(layer); + if (config.isFullAttentionLayer(layer)) { + auto decoder = reg(layer_name, config, layer); + decoder.attention_.cache_layer_index_ = mla_index; + mla_layers_.push_back(std::move(decoder)); + layer_types_.push_back(0); + layer_dispatch_.push_back(mla_index++); + } else { + kda_layers_.push_back(reg(layer_name, config, layer)); + layer_types_.push_back(1); + layer_dispatch_.push_back(kda_index++); + } + } + norm_ = reg("norm", config.rms_norm_eps, false); + } + + void resetKDAStates(int32_t batch_size) { + for (auto& layer : kda_layers_) { layer.attention_.resetState(batch_size); } + } + + std::vector forward(const std::vector& inputs, const std::vector& args) override { + auto hidden = word_embeddings_(inputs[0]); + for (size_t layer = 0; layer < layer_types_.size(); ++layer) { + if (layer_types_[layer] == 0) { + hidden = mla_layers_[layer_dispatch_[layer]](hidden, inputs[1], inputs[2], args[0])[0]; + } else { + hidden = kda_layers_[layer_dispatch_[layer]](hidden)[0]; + } + } + return {norm_(hidden)}; + } + + private: + nn::Embedding word_embeddings_; + nn::RMSNorm norm_; + std::vector mla_layers_; + std::vector kda_layers_; + std::vector layer_types_; + std::vector layer_dispatch_; +}; + +class Ling3ForCausalLM final : public ARGeneration, public nn::Module { + public: + explicit Ling3ForCausalLM(const Ling3Config& config) { + if (!hasOfficialLing3TinyArchitecture(config)) { + throw std::invalid_argument("Ling-3 CPU currently supports inclusionAI/Ling-3.0-tiny only"); + } + kv_cache_ = nn::StaticCache(config.max_cache_length, config.numFullAttentionLayers(), config.num_attention_heads, + config.num_key_value_heads, config.qk_head_dim, kFloat32, kFloat32, kCPU, false); + model_ = reg("model", config); + lm_head_ = reg("lm_head", config.hidden_size, config.vocab_size, false, config.linear_impl_type); + registerBuffer("inv_freq", makeLing3RoPEInvFreq(config.qk_rope_head_dim, config.rope_theta)); + eos_token_id_ = config.eos_token_id; + additional_eos_token_ids_.insert(config.end_of_text_token_id); + max_length_ = config.max_cache_length; + } + + ARGenerationOutputPast forward(const ARGenerationOutputPast& input, const ARGenerationArgs& args) override { + auto sequence = input.at("sequence"); + if (sequence.shape().size() != 2 || sequence.shape()[0] != 1 || sequence.shape()[1] <= 0 || sequence.dtype() != kInt64 + || sequence.device() != kCPU) { + throw std::invalid_argument("Ling-3 CPU expects a non-empty rank-2 int64 CPU sequence with batch size 1"); + } + const int sequence_length = sequence.shape()[1]; + const int cached_tokens = kv_cache_.getCurrentSeqCnt(0); + if (sequence_length > max_length_ || cached_tokens > max_length_ - sequence_length) { + throw std::invalid_argument("Ling-3 sequence exceeds the configured KV-cache capacity"); + } + + // Autoregressive input orchestration; position ids feed the model-local + // checkpoint-layout RoPE preparation above. + auto position_ids = Tensor::empty({1, sequence_length}, kInt64, kCPU).alloc(); + auto* positions = position_ids.ptr(); + if (input.count("position_ids") != 0 && sequence_length == 1) { + auto previous_positions = input.at("position_ids"); + previous_positions = previous_positions.contiguous(); + positions[0] = previous_positions.ptr()[previous_positions.numel() - 1] + 1; + } else { + for (int index = 0; index < sequence_length; ++index) { positions[index] = cached_tokens + index; } + } + auto [cos, sin] = makeLing3RotaryEmbedding(position_ids, getBuffer("inv_freq")); + auto hidden = model_(sequence, cos, sin, AnyValue(&kv_cache_))[0]; + hidden = hidden[{kAll, {sequence_length - 1}, kAll}]; + return {{"sequence", lm_head_(hidden)}, {"position_ids", position_ids}}; + } + + void resetState(int32_t batch_size = 1) { + if (batch_size != 1) { throw std::invalid_argument("Ling-3 CPU currently supports batch size 1 only"); } + kv_cache_.clearCache(); + model_.resetKDAStates(batch_size); + } + + nn::StaticCache& kvCache() { return kv_cache_; } + + private: + Ling3Model model_; + nn::Linear lm_head_; + nn::StaticCache kv_cache_; +}; + +} // namespace mllm::models::ling3 diff --git a/mllm/models/ling3/tokenization_ling3.hpp b/mllm/models/ling3/tokenization_ling3.hpp new file mode 100644 index 000000000..03d77bfb4 --- /dev/null +++ b/mllm/models/ling3/tokenization_ling3.hpp @@ -0,0 +1,314 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include "mllm/models/ARGeneration.hpp" +#include "mllm/preprocessor/tokenizers/AutoTokenizer.hpp" +#include "mllm/preprocessor/tokenizers/BPE.hpp" +#include "mllm/preprocessor/tokenizers/Unicode.hpp" +#include "mllm/preprocessor/tokenizers/llama_cpp_unicode/unicode.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mllm::models::ling3 { + +inline std::string normalizeLing3NFC(const std::string& text) { + const auto codepoints = unicode_cpts_from_utf8(text); + std::vector normalized; + normalized.reserve(codepoints.size()); + for (size_t index = 0; index < codepoints.size();) { + // Unicode Hangul composition is algorithmic and the bundled llama.cpp + // table does not represent it as an ordinary combining-mark pair. + constexpr uint32_t kLBase = 0x1100; + constexpr uint32_t kVBase = 0x1161; + constexpr uint32_t kTBase = 0x11A7; + constexpr uint32_t kSBase = 0xAC00; + constexpr uint32_t kLCount = 19; + constexpr uint32_t kVCount = 21; + constexpr uint32_t kTCount = 28; + const uint32_t current = codepoints[index]; + if (current >= kLBase && current < kLBase + kLCount && index + 1 < codepoints.size() && codepoints[index + 1] >= kVBase + && codepoints[index + 1] < kVBase + kVCount) { + const uint32_t l_index = current - kLBase; + const uint32_t v_index = codepoints[index + 1] - kVBase; + uint32_t syllable = kSBase + (l_index * kVCount + v_index) * kTCount; + index += 2; + if (index < codepoints.size() && codepoints[index] > kTBase && codepoints[index] < kTBase + kTCount) { + syllable += codepoints[index] - kTBase; + ++index; + } + normalized.push_back(syllable); + continue; + } + + // Preserve code points that are already NFC. The bundled normalizer's + // legacy NFD table stores only the base of a precomposed character, so + // normalizing an entire string would incorrectly turn "é" into "e". + // Restrict it to explicit base + combining-mark clusters. + size_t cluster_end = index + 1; + while (cluster_end < codepoints.size() && unicode_cpt_flags(codepoints[cluster_end]).is_accent_mark) { ++cluster_end; } + if (cluster_end > index + 1) { + const std::vector cluster(codepoints.begin() + static_cast(index), + codepoints.begin() + static_cast(cluster_end)); + const auto composed = unicode_cpts_normalize_nfc(cluster); + normalized.insert(normalized.end(), composed.begin(), composed.end()); + } else { + normalized.push_back(current); + } + index = cluster_end; + } + std::string result; + for (const auto codepoint : normalized) { result += unicode_cpt_to_utf8(codepoint); } + return result; +} + +inline bool ling3TokenizerMatchPattern(const std::wstring& text, size_t& position, std::wstring& matched) { + if (position >= text.size()) { return false; } + + static const std::wstring contractions[] = {L"'s", L"'d", L"'m", L"'t", L"'ll", L"'ve", L"'re"}; + for (const auto& contraction : contractions) { + bool matches = position + contraction.size() <= text.size(); + for (size_t index = 0; matches && index < contraction.size(); ++index) { + matches = std::towlower(text[position + index]) == contraction[index]; + } + if (matches) { + matched = text.substr(position, contraction.size()); + position += contraction.size(); + return true; + } + } + + // [^\r\n\p{L}\p{N}]?+\p{L}+ + { + const size_t original_position = position; + if (!preprocessor::isLetter(text[position]) && !preprocessor::isDigit(text[position]) && text[position] != L'\r' + && text[position] != L'\n') { + ++position; + } + if (position < text.size() && preprocessor::isLetter(text[position])) { + while (position < text.size() && preprocessor::isLetter(text[position])) { ++position; } + matched = text.substr(original_position, position - original_position); + return true; + } + position = original_position; + } + + // \p{N} + if (preprocessor::isDigit(text[position])) { + matched = text.substr(position, 1); + ++position; + return true; + } + + // " ?[^\s\p{L}\p{N}]++[\r\n]*" + { + const size_t original_position = position; + if (text[position] == L' ') { ++position; } + if (position < text.size() && !std::iswspace(text[position]) && !preprocessor::isLetter(text[position]) + && !preprocessor::isDigit(text[position])) { + while (position < text.size() && !std::iswspace(text[position]) && !preprocessor::isLetter(text[position]) + && !preprocessor::isDigit(text[position])) { + ++position; + } + while (position < text.size() && (text[position] == L'\r' || text[position] == L'\n')) { ++position; } + matched = text.substr(original_position, position - original_position); + return true; + } + position = original_position; + } + + // \s*[\r\n]. Greedy \s* backtracks to the last available line break. + { + const size_t start = position; + size_t scan = position; + size_t last_line_break = std::wstring::npos; + while (scan < text.size() && std::iswspace(text[scan])) { + if (text[scan] == L'\r' || text[scan] == L'\n') { last_line_break = scan + 1; } + ++scan; + } + if (last_line_break != std::wstring::npos) { + position = last_line_break; + matched = text.substr(start, position - start); + return true; + } + } + + // \s+(?!\S), followed by the final \s+ fallback. + if (std::iswspace(text[position])) { + const size_t start = position; + while (position < text.size() && std::iswspace(text[position])) { ++position; } + if (position >= text.size()) { + matched = text.substr(start, position - start); + return true; + } + if (position - start > 1) { + --position; + matched = text.substr(start, position - start); + return true; + } + position = start; + while (position < text.size() && std::iswspace(text[position])) { ++position; } + matched = text.substr(start, position - start); + return true; + } + return false; +} + +inline std::vector ling3RegexSplit(const std::string& text) { + const auto wide_text = preprocessor::utf8string2WideString(text); + std::vector pieces; + size_t position = 0; + while (position < wide_text.size()) { + std::wstring matched; + if (ling3TokenizerMatchPattern(wide_text, position, matched)) { + pieces.push_back(std::move(matched)); + } else { + pieces.push_back(wide_text.substr(position, 1)); + ++position; + } + } + return pieces; +} + +struct Ling3Message { + std::string prompt; + std::string system_prompt; + bool enable_thinking = true; +}; + +class Ling3Tokenizer final : public preprocessor::AutoTokenizer { + public: + explicit Ling3Tokenizer(const std::string& file_path) { + preprocessor::initLocal(); + preprocessor::makeBytes2UnicodeMap(bytes_to_unicode_); + for (const auto& [byte, codepoint] : bytes_to_unicode_) { unicode_to_bytes_.insert({codepoint, byte}); } + if (!bpe_.initFromSentencePieceJson(file_path)) { + throw std::invalid_argument("Ling-3 tokenizer could not load tokenizer.json"); + } + + std::ifstream tokenizer_file(file_path); + if (!tokenizer_file.is_open()) { throw std::invalid_argument("Ling-3 tokenizer.json is not readable"); } + const auto tokenizer_json = nlohmann::json::parse(tokenizer_file); + if (!tokenizer_json.contains("normalizer") || tokenizer_json["normalizer"].value("type", "") != "NFC") { + throw std::invalid_argument("Ling-3 tokenizer requires the official NFC normalizer"); + } + if (!tokenizer_json.contains("added_tokens")) { + throw std::invalid_argument("Ling-3 tokenizer.json is missing added_tokens"); + } + for (const auto& added_token : tokenizer_json["added_tokens"]) { + const auto content = added_token.at("content").get(); + const auto wide_content = preprocessor::utf8string2WideString(content); + special_tokens_trie_.add(wide_content); + added_tokens_.insert(wide_content); + } + } + + std::vector _tokenize(const std::string& text) override { + std::vector tokens; + for (const auto& piece : ling3RegexSplit(normalizeLing3NFC(text))) { + const auto utf8_piece = preprocessor::wideString2Utf8String(piece); + std::wstring mapped; + for (const unsigned char byte : utf8_piece) { mapped.push_back(bytes_to_unicode_.at(byte)); } + auto bpe_tokens = bpe_._bpe(mapped); + tokens.insert(tokens.end(), bpe_tokens.begin(), bpe_tokens.end()); + } + return tokens; + } + + std::vector tokenize(const std::string& text) override { + const auto split_tokens = special_tokens_trie_.split(preprocessor::utf8string2WideString(text)); + std::vector tokens; + for (const auto& token : split_tokens) { + if (special_tokens_trie_.isSpecialToken(token)) { + tokens.push_back(token); + } else { + auto ordinary_tokens = _tokenize(preprocessor::wideString2Utf8String(token)); + tokens.insert(tokens.end(), ordinary_tokens.begin(), ordinary_tokens.end()); + } + } + return tokens; + } + + std::wstring _detokenize(int64_t token_id) override { return bpe_._lookup_inverse_vocab(token_id); } + + std::string detokenizeBytes(int64_t token_id) { + const auto token = _detokenize(token_id); + if (added_tokens_.count(token) != 0) { return preprocessor::wideString2Utf8String(token); } + std::string bytes; + for (const auto codepoint : token) { + const auto iterator = unicode_to_bytes_.find(codepoint); + if (iterator == unicode_to_bytes_.end()) { + throw std::runtime_error("Ling-3 tokenizer encountered an unknown byte-unicode symbol"); + } + bytes.push_back(static_cast(iterator->second)); + } + return bytes; + } + + std::wstring detokenize(int64_t token_id) override { return preprocessor::utf8string2WideString(detokenizeBytes(token_id)); } + + Tensor convert2Ids(const std::vector& tokens) override { + auto result = Tensor::empty({1, static_cast(tokens.size())}, kInt64, kCPU) + .setMemType(kExtraInput) + .setName("ling3-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; + } + + ARGenerationOutputPast convertMessage(const Ling3Message& message) { + if (message.prompt.empty()) { throw std::invalid_argument("Ling-3 prompt must not be empty"); } + static constexpr std::string_view kReservedMarkers[] = { + "SYSTEM", + "HUMAN", + "ASSISTANT", + "<|role_end|>", + }; + for (const auto marker : kReservedMarkers) { + if (message.prompt.find(marker) != std::string::npos) { + throw std::invalid_argument("Ling-3 prompt must not inject reserved chat-template markers"); + } + } + + const std::string thinking_option = message.enable_thinking ? "on" : "off"; + std::string rendered = "SYSTEM"; + if (message.system_prompt.empty()) { + rendered += "detailed thinking " + thinking_option + "<|role_end|>"; + } else if (message.system_prompt.find("detailed thinking on") != std::string::npos + || message.system_prompt.find("detailed thinking off") != std::string::npos) { + rendered += message.system_prompt + "<|role_end|>"; + } else { + rendered += message.system_prompt + "\ndetailed thinking " + thinking_option + "<|role_end|>"; + } + rendered += "HUMAN" + message.prompt + "<|role_end|>ASSISTANT"; + rendered += message.enable_thinking ? "\n" : "\n"; + + const auto tokens = tokenize(rendered); + auto sequence = Tensor::empty({1, static_cast(tokens.size())}, kInt64, kCPU) + .setMemType(kNormal) + .setName("ling3-tokenizer-i0") + .alloc(); + auto* ids = sequence.ptr(); + for (size_t index = 0; index < tokens.size(); ++index) { ids[index] = bpe_._lookup_vocab(tokens[index]); } + return {{"sequence", sequence}}; + } + + private: + preprocessor::BPE bpe_; + std::unordered_map bytes_to_unicode_; + std::unordered_map unicode_to_bytes_; + std::unordered_set added_tokens_; +}; + +} // namespace mllm::models::ling3 diff --git a/mllm/nn/Functional.cpp b/mllm/nn/Functional.cpp index 5f78c13eb..7f8bd5ce3 100644 --- a/mllm/nn/Functional.cpp +++ b/mllm/nn/Functional.cpp @@ -27,6 +27,7 @@ #include "mllm/core/aops/RadixAttnWithSinkAndSwaDiffDimOp.hpp" #include "mllm/core/aops/WhereOp.hpp" #include "mllm/core/aops/GatedDeltaRuleOp.hpp" +#include "mllm/core/aops/KimiDeltaAttentionOp.hpp" #include "mllm/engine/Context.hpp" namespace mllm::nn::functional { @@ -250,6 +251,16 @@ mllm::Tensor sigmoid(const Tensor& x) { return ctx.buildOpAndSubmitTask(OpTypes::kSigmoid, aops::SigmoidOpOptions{}, {x})[0]; } +std::array kimiDeltaAttention(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& gate_logits, + const Tensor& beta, const Tensor& a_log, const Tensor& dt_bias, const Tensor& state, + bool safe_gate, float lower_bound, bool state_inplace) { + auto outputs = Context::instance().buildOpAndSubmitTask( + OpTypes::kKimiDeltaAttention, + aops::KimiDeltaAttentionOpOptions{.safe_gate = safe_gate, .lower_bound = lower_bound, .state_inplace = state_inplace}, + {q, k, v, gate_logits, beta, a_log, dt_bias, state}); + return {outputs[0], outputs[1]}; +} + mllm::Tensor gather(const Tensor& x, int dim, const Tensor& indices) { auto& ctx = mllm::Context::instance(); return ctx.buildOpAndSubmitTask(OpTypes::kGather, aops::GatherOpOptions{.dim = dim}, {x, indices})[0]; diff --git a/mllm/nn/Functional.hpp b/mllm/nn/Functional.hpp index e7c1e135c..f75daac91 100644 --- a/mllm/nn/Functional.hpp +++ b/mllm/nn/Functional.hpp @@ -14,6 +14,7 @@ #include "mllm/core/aops/InterpolateOp.hpp" #include "mllm/core/aops/GroupedQueryAttentionOp.hpp" #include "mllm/core/aops/GatedDeltaRuleOp.hpp" +#include "mllm/core/aops/KimiDeltaAttentionOp.hpp" #include "mllm/core/aops/RadixAttnWithSinkAndSwaDiffDimOp.hpp" #include "mllm/engine/Context.hpp" @@ -178,6 +179,10 @@ mllm::Tensor where(const Tensor& mask, const Tensor& original, const Tensor& v); mllm::Tensor sigmoid(const Tensor& x); +std::array kimiDeltaAttention(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& gate_logits, + const Tensor& beta, const Tensor& a_log, const Tensor& dt_bias, const Tensor& state, + bool safe_gate = true, float lower_bound = -5.0F, bool state_inplace = false); + mllm::Tensor gather(const Tensor& x, int dim, const Tensor& indices); } // namespace mllm::nn::functional diff --git a/mllm/nn/Nn.hpp b/mllm/nn/Nn.hpp index 3b0ab764a..a05181700 100644 --- a/mllm/nn/Nn.hpp +++ b/mllm/nn/Nn.hpp @@ -36,3 +36,4 @@ #include "mllm/nn/layers/LayerNorm2D.hpp" // IWYU pragma: export #include "mllm/nn/layers/CausalDepthwiseConv1D.hpp" // IWYU pragma: export #include "mllm/nn/layers/GatedDeltaRule.hpp" // IWYU pragma: export +#include "mllm/nn/layers/KimiDeltaAttention.hpp" // IWYU pragma: export diff --git a/mllm/nn/layers/KimiDeltaAttention.cpp b/mllm/nn/layers/KimiDeltaAttention.cpp new file mode 100644 index 000000000..931c20d3b --- /dev/null +++ b/mllm/nn/layers/KimiDeltaAttention.cpp @@ -0,0 +1,18 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/nn/layers/KimiDeltaAttention.hpp" + +namespace mllm::nn { + +KimiDeltaAttention::KimiDeltaAttention() : Layer(OpTypes::kKimiDeltaAttention, aops::KimiDeltaAttentionOpOptions{}) {} + +KimiDeltaAttention::KimiDeltaAttention(const aops::KimiDeltaAttentionOpOptions& options) + : Layer(OpTypes::kKimiDeltaAttention, options) {} + +KimiDeltaAttention::KimiDeltaAttention(bool safe_gate, float lower_bound, bool state_inplace) + : Layer(OpTypes::kKimiDeltaAttention, + aops::KimiDeltaAttentionOpOptions{ + .safe_gate = safe_gate, .lower_bound = lower_bound, .state_inplace = state_inplace}) {} + +} // namespace mllm::nn diff --git a/mllm/nn/layers/KimiDeltaAttention.hpp b/mllm/nn/layers/KimiDeltaAttention.hpp new file mode 100644 index 000000000..1e25fa972 --- /dev/null +++ b/mllm/nn/layers/KimiDeltaAttention.hpp @@ -0,0 +1,20 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include "mllm/core/aops/KimiDeltaAttentionOp.hpp" +#include "mllm/nn/Layer.hpp" + +namespace mllm::nn { + +class KimiDeltaAttention : public Layer { + public: + KimiDeltaAttention(); + explicit KimiDeltaAttention(const aops::KimiDeltaAttentionOpOptions& options); + KimiDeltaAttention(bool safe_gate, float lower_bound, bool state_inplace = false); + + MLLM_LAYER_ANY_INPUTS_2_OUTPUTS_FORWARD +}; + +} // namespace mllm::nn diff --git a/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index 70ba14ac4..027f16df2 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -31,5 +31,6 @@ add_test( NAME CPUKernelFocused COMMAND Mllm-Test-CPUKernel - --gtest_filter=CausalDepthwiseConvKernelTest.*:CausalDepthwiseConvCurrentFirstKernelTest.*:GatedDeltaRuleKernelTest.*) + --gtest_filter=CausalDepthwiseConvKernelTest.*:CausalDepthwiseConvCurrentFirstKernelTest.*:GatedDeltaRuleKernelTest.*:KimiDeltaAttentionKernelTest.* +) set_tests_properties(CPUKernelFocused PROPERTIES LABELS cpu-kernel) diff --git a/tests/cpu/KaiW4A32PackTest.cpp b/tests/cpu/KaiW4A32PackTest.cpp index 78653639e..f1653d022 100644 --- a/tests/cpu/KaiW4A32PackTest.cpp +++ b/tests/cpu/KaiW4A32PackTest.cpp @@ -14,6 +14,7 @@ #include "mllm/backends/cpu/kernels/common/kai_w4a32_pack.hpp" #if defined(__aarch64__) +#include "mllm/backends/cpu/kernels/arm/linear/kai.hpp" #include "kai_matmul_clamp_f32_qai8dxp1x4_qsi4c32p4x4_1x4_neon_dotprod.h" #include "kai_matmul_clamp_f32_qai8dxp1x8_qsi4c32p4x8_1x4x32_neon_dotprod.h" #include "kai_matmul_clamp_f32_qai8dxp1x8_qsi4c32p8x8_1x8x32_neon_dotprod.h" @@ -91,6 +92,44 @@ TEST(KaiW4A32PackTest, SharedTileParametersMatchArmUkernels) { EXPECT_EQ(tile.sr, test_case.get_sr()) << test_case.name; } } + +TEST(KaiW4A32PackTest, BatchedRowsMatchIndependentRows) { + constexpr int kRows = 7; + constexpr int kOutChannels = 17; + constexpr int kInChannels = 64; + using Helper = mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk; + constexpr auto kTile = Helper::Tiles::qai8dxp1x8_qsi4c32p8x8_1x8x32; + + std::vector weights(kOutChannels * kInChannels); + std::vector input(kRows * kInChannels); + for (int index = 0; index < static_cast(weights.size()); ++index) { + weights[index] = std::sin(static_cast(index) * 0.071F) * 0.30F + - std::cos(static_cast(index) * 0.019F) * 0.05F; + } + for (int index = 0; index < static_cast(input.size()); ++index) { + input[index] = std::sin(static_cast(index) * 0.037F) * 0.50F + + std::cos(static_cast(index) * 0.013F) * 0.07F; + } + + Helper helper; + std::vector packed(helper.quant_pack_rhs_size(kOutChannels, kInChannels, kTile)); + helper.quant_pack_rhs_offline(packed.data(), weights.data(), nullptr, kOutChannels, kInChannels, kTile); + std::vector batched_workspace(helper.workspace_size(kRows, kInChannels, kTile)); + std::vector batched(kRows * kOutChannels); + helper.matmul(batched.data(), input.data(), packed.data(), batched_workspace.data(), kRows, kInChannels, kOutChannels, + kTile, 1); + + std::vector row_workspace(helper.workspace_size(1, kInChannels, kTile)); + std::vector row(kOutChannels); + for (int row_index = 0; row_index < kRows; ++row_index) { + helper.matmul(row.data(), input.data() + row_index * kInChannels, packed.data(), row_workspace.data(), 1, kInChannels, + kOutChannels, kTile, 1); + for (int channel = 0; channel < kOutChannels; ++channel) { + EXPECT_FLOAT_EQ(batched[row_index * kOutChannels + channel], row[channel]) + << "row=" << row_index << " channel=" << channel; + } + } +} #endif } // namespace diff --git a/tests/cpu/KernelTest.cpp b/tests/cpu/KernelTest.cpp index b3f3d4d99..9e1a6892f 100644 --- a/tests/cpu/KernelTest.cpp +++ b/tests/cpu/KernelTest.cpp @@ -630,6 +630,38 @@ TEST_F(GatedDeltaRuleKernelTest, ProductionGroupedHeadGeometry8LaneIsBitwiseStab /*repeats=*/24)); } +//===----------------------------------------------------------------------===// +// Kimi delta attention +//===----------------------------------------------------------------------===// +#include "KimiDeltaAttentionKernelTest.hpp" +TEST_F(KimiDeltaAttentionKernelTest, MatchesScalarReferenceAcrossGateVariants) { + EXPECT_NO_FATAL_FAILURE(testMatchesScalarReference({ + // Scalar tails only (D < 4) and exact NEON lane blocks, both gates. + {/*batch=*/1, /*sequence=*/1, /*heads=*/1, /*dim=*/3, /*safe_gate=*/true, /*lower_bound=*/-5.0F, /*threads=*/1}, + {/*batch=*/2, /*sequence=*/5, /*heads=*/3, /*dim=*/8, /*safe_gate=*/true, /*lower_bound=*/-5.0F, /*threads=*/1}, + {/*batch=*/2, /*sequence=*/5, /*heads=*/3, /*dim=*/8, /*safe_gate=*/false, /*lower_bound=*/0.0F, /*threads=*/1}, + {/*batch=*/1, /*sequence=*/4, /*heads=*/2, /*dim=*/6, /*safe_gate=*/true, /*lower_bound=*/-2.5F, /*threads=*/1}, + // Ling-3.0-tiny head geometry with the multi-threaded (batch, head) partition. + {/*batch=*/1, /*sequence=*/9, /*heads=*/16, /*dim=*/128, /*safe_gate=*/true, /*lower_bound=*/-5.0F, /*threads=*/8}, + })); +} + +TEST_F(KimiDeltaAttentionKernelTest, PrefillAndTokenwiseDecodeAreBitwiseEqual) { + EXPECT_NO_FATAL_FAILURE(testPrefillAndTokenwiseDecodeAreBitwiseEqual( + {/*batch=*/1, /*sequence=*/7, /*heads=*/2, /*dim=*/8, /*safe_gate=*/true, /*lower_bound=*/-5.0F, /*threads=*/1})); + EXPECT_NO_FATAL_FAILURE(testPrefillAndTokenwiseDecodeAreBitwiseEqual( + {/*batch=*/1, /*sequence=*/49, /*heads=*/16, /*dim=*/128, /*safe_gate=*/true, /*lower_bound=*/-5.0F, /*threads=*/8})); +} + +TEST_F(KimiDeltaAttentionKernelTest, ParallelLanesMatchSerialBitwise) { + EXPECT_NO_FATAL_FAILURE(testParallelLanesMatchSerialBitwise( + {/*batch=*/2, /*sequence=*/16, /*heads=*/16, /*dim=*/128, /*safe_gate=*/true, /*lower_bound=*/-5.0F, /*threads=*/8})); +} + +TEST_F(KimiDeltaAttentionKernelTest, RejectsNullBuffersAndInvalidGeometry) { + EXPECT_NO_FATAL_FAILURE(testRejectsNullBuffersAndInvalidGeometry()); +} + //===----------------------------------------------------------------------===// // Parallel linear //===----------------------------------------------------------------------===// diff --git a/tests/cpu/KimiDeltaAttentionKernelTest.hpp b/tests/cpu/KimiDeltaAttentionKernelTest.hpp new file mode 100644 index 000000000..b359cfa05 --- /dev/null +++ b/tests/cpu/KimiDeltaAttentionKernelTest.hpp @@ -0,0 +1,249 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +// Focused oracle for the Kimi Delta Attention (KDA) recurrent kernel. +// +// The reference below is an independent scalar implementation of the frozen +// contract (L2-normalised q/k, 1/sqrt(D) query scaling, safe-gate or softplus +// log-decay, delta-rule state update). It is deliberately not routed through +// the production kernel, so the NEON row helpers inside kimiDeltaAttentionF32 +// cannot validate themselves. Output and final state are compared separately: +// an output-only comparison would miss a corrupted state that only shows up in +// the next chunk. + +#include + +#include +#include +#include +#include + +#include "mllm/backends/cpu/kernels/common/kda/kimi_delta_attention.hpp" +#include "KernelTestHelper.hpp" + +namespace kimi_delta_attention_kernel_test { + +using mllm::cpu::kda::kimiDeltaAttentionF32; + +struct Geometry { + int batch = 1; + int sequence = 1; + int heads = 1; + int dim = 1; + bool safe_gate = true; + float lower_bound = -5.0F; + int thread_count = 1; + + [[nodiscard]] std::size_t vectorCount() const { + return static_cast(batch) * sequence * heads * dim; + } + [[nodiscard]] std::size_t betaCount() const { return static_cast(batch) * sequence * heads; } + [[nodiscard]] std::size_t stateCount() const { return static_cast(batch) * heads * dim * dim; } +}; + +// Deterministic index-derived fill. No RNG, so every host reproduces the same +// bytes without carrying a seed through the evidence record. +inline float patternValue(std::size_t index, float scale, float offset) { + return offset + scale * static_cast(static_cast(index * 37U % 29U) - 14); +} + +inline std::vector makeBuffer(std::size_t count, float scale, float offset = 0.0F) { + std::vector values(count); + for (std::size_t index = 0; index < count; ++index) { values[index] = patternValue(index, scale, offset); } + return values; +} + +struct Inputs { + std::vector q, k, v, gate, beta, a_log, dt_bias, state; +}; + +inline Inputs makeInputs(const Geometry& geometry, float scale = 1.0F) { + Inputs inputs; + inputs.q = makeBuffer(geometry.vectorCount(), 0.017F * scale); + inputs.k = makeBuffer(geometry.vectorCount(), -0.013F * scale, 0.02F * scale); + inputs.v = makeBuffer(geometry.vectorCount(), 0.011F * scale, -0.03F * scale); + inputs.gate = makeBuffer(geometry.vectorCount(), 0.019F * scale, 0.1F * scale); + inputs.beta = makeBuffer(geometry.betaCount(), 0.01F * scale, 0.45F); + inputs.a_log = makeBuffer(static_cast(geometry.heads), 0.02F * scale, 0.3F); + inputs.dt_bias = makeBuffer(static_cast(geometry.heads) * geometry.dim, 0.015F * scale, -0.1F * scale); + inputs.state = makeBuffer(geometry.stateCount(), 0.003F * scale); + return inputs; +} + +inline float referenceSigmoid(float value) { return 1.0F / (1.0F + std::exp(-value)); } + +// Independent scalar reference for the frozen contract: +// q/k/v/gate [B, S, H, D], beta [B, S, H], a_log [H], dt_bias [H, D], +// state [B, H, D, D] updated in place, output [B, S, H, D]. +inline void referenceKimiDeltaAttention(const Inputs& inputs, std::vector& state, std::vector& output, + const Geometry& geometry) { + const int dim = geometry.dim; + const float query_dim_scale = 1.0F / std::sqrt(static_cast(dim)); + for (int b = 0; b < geometry.batch; ++b) { + for (int h = 0; h < geometry.heads; ++h) { + const std::size_t state_base = (static_cast(b) * geometry.heads + h) * dim * dim; + for (int s = 0; s < geometry.sequence; ++s) { + const std::size_t vector_base = ((static_cast(b) * geometry.sequence + s) * geometry.heads + h) * dim; + float q_norm_sq = 0.0F; + float k_norm_sq = 0.0F; + for (int d = 0; d < dim; ++d) { + q_norm_sq += inputs.q[vector_base + d] * inputs.q[vector_base + d]; + k_norm_sq += inputs.k[vector_base + d] * inputs.k[vector_base + d]; + } + const float q_scale = query_dim_scale / std::sqrt(q_norm_sq + 1.0e-6F); + const float k_scale = 1.0F / std::sqrt(k_norm_sq + 1.0e-6F); + std::vector normalized_q(dim); + std::vector normalized_k(dim); + std::vector prediction(dim, 0.0F); + for (int d = 0; d < dim; ++d) { + normalized_q[d] = inputs.q[vector_base + d] * q_scale; + normalized_k[d] = inputs.k[vector_base + d] * k_scale; + } + for (int kd = 0; kd < dim; ++kd) { + const float gate_input = inputs.gate[vector_base + kd] + inputs.dt_bias[h * dim + kd]; + const float log_decay = geometry.safe_gate + ? geometry.lower_bound * referenceSigmoid(std::exp(inputs.a_log[h]) * gate_input) + : -std::exp(inputs.a_log[h]) * std::log1p(std::exp(gate_input)); + const float decay = std::exp(log_decay); + for (int vd = 0; vd < dim; ++vd) { + const std::size_t state_index = state_base + kd * dim + vd; + state[state_index] *= decay; + prediction[vd] += normalized_k[kd] * state[state_index]; + } + } + const float beta = inputs.beta[(static_cast(b) * geometry.sequence + s) * geometry.heads + h]; + for (int vd = 0; vd < dim; ++vd) { + const float delta = beta * (inputs.v[vector_base + vd] - prediction[vd]); + for (int kd = 0; kd < dim; ++kd) { + const std::size_t state_index = state_base + kd * dim + vd; + state[state_index] += normalized_k[kd] * delta; + output[vector_base + vd] += normalized_q[kd] * state[state_index]; + } + } + } + } + } +} + +// Output and updated state must both match the scalar reference within a +// small tolerance: the NEON helpers reorder the D-lane accumulations. +inline void testMatchesScalarReference(const std::vector& geometries, float tolerance = 2.0e-6F) { + for (const auto& geometry : geometries) { + SCOPED_TRACE(::testing::Message() << "B=" << geometry.batch << " S=" << geometry.sequence << " H=" << geometry.heads + << " D=" << geometry.dim << " safe_gate=" << geometry.safe_gate + << " threads=" << geometry.thread_count); + const auto inputs = makeInputs(geometry); + auto expected_state = inputs.state; + auto actual_state = inputs.state; + std::vector expected_output(geometry.vectorCount(), 0.0F); + std::vector actual_output(geometry.vectorCount(), 0.0F); + + referenceKimiDeltaAttention(inputs, expected_state, expected_output, geometry); + kimiDeltaAttentionF32(inputs.q.data(), inputs.k.data(), inputs.v.data(), inputs.gate.data(), inputs.beta.data(), + inputs.a_log.data(), inputs.dt_bias.data(), actual_state.data(), actual_output.data(), + geometry.batch, geometry.sequence, geometry.heads, geometry.dim, geometry.safe_gate, + geometry.lower_bound, geometry.thread_count); + + for (std::size_t i = 0; i < expected_output.size(); ++i) { + ASSERT_NEAR(actual_output[i], expected_output[i], tolerance) << "output index " << i; + } + for (std::size_t i = 0; i < expected_state.size(); ++i) { + ASSERT_NEAR(actual_state[i], expected_state[i], tolerance) << "state index " << i; + } + } +} + +// One-shot prefill and token-by-token decode must be bitwise identical: the +// recurrence has no cross-token reduction, so any drift would be a state bug. +inline void testPrefillAndTokenwiseDecodeAreBitwiseEqual(const Geometry& geometry) { + ASSERT_EQ(geometry.batch, 1) << "tokenwise replay is defined for batch 1"; + const auto inputs = makeInputs(geometry); + std::vector prefill_state(geometry.stateCount(), 0.0F); + std::vector decode_state(geometry.stateCount(), 0.0F); + std::vector prefill_output(geometry.vectorCount(), 0.0F); + std::vector decode_output(geometry.vectorCount(), 0.0F); + + kimiDeltaAttentionF32(inputs.q.data(), inputs.k.data(), inputs.v.data(), inputs.gate.data(), inputs.beta.data(), + inputs.a_log.data(), inputs.dt_bias.data(), prefill_state.data(), prefill_output.data(), + geometry.batch, geometry.sequence, geometry.heads, geometry.dim, geometry.safe_gate, + geometry.lower_bound, geometry.thread_count); + const std::size_t token_width = static_cast(geometry.heads) * geometry.dim; + for (int token = 0; token < geometry.sequence; ++token) { + const std::size_t offset = static_cast(token) * token_width; + kimiDeltaAttentionF32(inputs.q.data() + offset, inputs.k.data() + offset, inputs.v.data() + offset, + inputs.gate.data() + offset, inputs.beta.data() + static_cast(token) * geometry.heads, + inputs.a_log.data(), inputs.dt_bias.data(), decode_state.data(), decode_output.data() + offset, + geometry.batch, 1, geometry.heads, geometry.dim, geometry.safe_gate, geometry.lower_bound, + geometry.thread_count); + } + + EXPECT_EQ(prefill_output, decode_output); + EXPECT_EQ(prefill_state, decode_state); +} + +// Lanes partition (batch, head) tasks and never share state rows, so a +// multi-threaded run must reproduce the serial run bitwise. +inline void testParallelLanesMatchSerialBitwise(const Geometry& geometry) { + ASSERT_GT(geometry.thread_count, 1); + const auto inputs = makeInputs(geometry, 0.1F); + auto serial_state = inputs.state; + auto parallel_state = inputs.state; + std::vector serial_output(geometry.vectorCount(), 0.0F); + std::vector parallel_output(geometry.vectorCount(), 0.0F); + + kimiDeltaAttentionF32(inputs.q.data(), inputs.k.data(), inputs.v.data(), inputs.gate.data(), inputs.beta.data(), + inputs.a_log.data(), inputs.dt_bias.data(), serial_state.data(), serial_output.data(), geometry.batch, + geometry.sequence, geometry.heads, geometry.dim, geometry.safe_gate, geometry.lower_bound, + /*thread_count=*/1); + kimiDeltaAttentionF32(inputs.q.data(), inputs.k.data(), inputs.v.data(), inputs.gate.data(), inputs.beta.data(), + inputs.a_log.data(), inputs.dt_bias.data(), parallel_state.data(), parallel_output.data(), + geometry.batch, geometry.sequence, geometry.heads, geometry.dim, geometry.safe_gate, + geometry.lower_bound, geometry.thread_count); + + EXPECT_EQ(serial_output, parallel_output); + EXPECT_EQ(serial_state, parallel_state); +} + +inline void testRejectsNullBuffersAndInvalidGeometry() { + std::vector values(8, 0.0F); + std::vector state(8, 0.0F); + const auto call = [&](const float* q, int batch, int sequence, int heads, int dim, bool safe_gate, float lower_bound, + int threads) { + kimiDeltaAttentionF32(q, values.data(), values.data(), values.data(), values.data(), values.data(), values.data(), + state.data(), values.data(), batch, sequence, heads, dim, safe_gate, lower_bound, threads); + }; + EXPECT_THROW(call(nullptr, 1, 1, 1, 2, true, -5.0F, 1), std::invalid_argument); + EXPECT_THROW(call(values.data(), 0, 1, 1, 2, true, -5.0F, 1), std::invalid_argument); + EXPECT_THROW(call(values.data(), 1, 1, 1, 2, true, -5.0F, 0), std::invalid_argument); + // The safe gate needs a finite, strictly negative lower bound. + EXPECT_THROW(call(values.data(), 1, 1, 1, 2, true, 0.0F, 1), std::invalid_argument); + EXPECT_THROW(call(values.data(), 1, 1, 1, 2, true, INFINITY, 1), std::invalid_argument); + // The softplus gate ignores the bound entirely. + EXPECT_NO_THROW(call(values.data(), 1, 1, 1, 2, false, 0.0F, 1)); +} + +} // namespace kimi_delta_attention_kernel_test + +class KimiDeltaAttentionKernelTest : public KernelTest { + public: + KimiDeltaAttentionKernelTest() = default; + ~KimiDeltaAttentionKernelTest() override = default; + + void testMatchesScalarReference(const std::vector& geometries) { + kimi_delta_attention_kernel_test::testMatchesScalarReference(geometries); + } + + void testPrefillAndTokenwiseDecodeAreBitwiseEqual(const kimi_delta_attention_kernel_test::Geometry& geometry) { + kimi_delta_attention_kernel_test::testPrefillAndTokenwiseDecodeAreBitwiseEqual(geometry); + } + + void testParallelLanesMatchSerialBitwise(const kimi_delta_attention_kernel_test::Geometry& geometry) { + kimi_delta_attention_kernel_test::testParallelLanesMatchSerialBitwise(geometry); + } + + void testRejectsNullBuffersAndInvalidGeometry() { + kimi_delta_attention_kernel_test::testRejectsNullBuffersAndInvalidGeometry(); + } +}; diff --git a/tests/cpu/Ling3ConfigTest.cpp b/tests/cpu/Ling3ConfigTest.cpp new file mode 100644 index 000000000..4311c1760 --- /dev/null +++ b/tests/cpu/Ling3ConfigTest.cpp @@ -0,0 +1,28 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/models/ling3/configuration_ling3.hpp" + +#include + +#include +#include + +namespace { + +std::string ling3RuntimeConfigPath() { + const char* value = std::getenv("LING3_RUNTIME_CONFIG"); + return value == nullptr ? std::string(LING3_EXAMPLE_DIR) + "/config_tiny_w4a32_kai.json" : std::string(value); +} + +} // namespace + +TEST(Ling3Config, AcceptsPinnedTinyRuntimeContract) { + const mllm::models::ling3::Ling3Config config(ling3RuntimeConfigPath()); + EXPECT_TRUE(mllm::models::ling3::hasOfficialLing3TinyArchitecture(config)); + EXPECT_EQ(config.numFullAttentionLayers(), 6); + EXPECT_EQ(config.numKDALayers(), 18); + for (int layer = 0; layer < config.num_hidden_layers; ++layer) { + EXPECT_EQ(config.isFullAttentionLayer(layer), (layer + 1) % 4 == 0); + } +} diff --git a/tests/cpu/Ling3KDATest.cpp b/tests/cpu/Ling3KDATest.cpp new file mode 100644 index 000000000..596423462 --- /dev/null +++ b/tests/cpu/Ling3KDATest.cpp @@ -0,0 +1,312 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/kernels/common/kda/kimi_delta_attention.hpp" +#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/KimiDeltaAttentionOp.hpp" +#include "mllm/mllm.hpp" +#include "mllm/nn/Functional.hpp" + +#include + +#include +#include +#include +#include + +#include + +namespace { + +float sigmoid(float value) { return 1.0F / (1.0F + std::exp(-value)); } + +void referenceKDA(const std::vector& q, const std::vector& k, const std::vector& v, + const std::vector& gate, const std::vector& beta, const std::vector& a_log, + const std::vector& dt_bias, std::vector& state, std::vector& output, int batch, + int sequence, int heads, int dim, bool safe_gate, float lower_bound) { + const float query_dim_scale = 1.0F / std::sqrt(static_cast(dim)); + for (int b = 0; b < batch; ++b) { + for (int h = 0; h < heads; ++h) { + const std::size_t state_base = (static_cast(b) * heads + h) * dim * dim; + for (int s = 0; s < sequence; ++s) { + const std::size_t vector_base = ((static_cast(b) * sequence + s) * heads + h) * dim; + float q_norm_sq = 0.0F; + float k_norm_sq = 0.0F; + for (int d = 0; d < dim; ++d) { + q_norm_sq += q[vector_base + d] * q[vector_base + d]; + k_norm_sq += k[vector_base + d] * k[vector_base + d]; + } + const float q_scale = query_dim_scale / std::sqrt(q_norm_sq + 1.0e-6F); + const float k_scale = 1.0F / std::sqrt(k_norm_sq + 1.0e-6F); + std::vector normalized_q(dim); + std::vector normalized_k(dim); + std::vector prediction(dim, 0.0F); + for (int d = 0; d < dim; ++d) { + normalized_q[d] = q[vector_base + d] * q_scale; + normalized_k[d] = k[vector_base + d] * k_scale; + } + for (int kd = 0; kd < dim; ++kd) { + const float gate_input = gate[vector_base + kd] + dt_bias[h * dim + kd]; + const float log_decay = safe_gate ? lower_bound * sigmoid(std::exp(a_log[h]) * gate_input) + : -std::exp(a_log[h]) * std::log1p(std::exp(gate_input)); + const float decay = std::exp(log_decay); + for (int vd = 0; vd < dim; ++vd) { + const std::size_t state_index = state_base + kd * dim + vd; + state[state_index] *= decay; + prediction[vd] += normalized_k[kd] * state[state_index]; + } + } + for (int vd = 0; vd < dim; ++vd) { + const float delta = + beta[(static_cast(b) * sequence + s) * heads + h] * (v[vector_base + vd] - prediction[vd]); + for (int kd = 0; kd < dim; ++kd) { + const std::size_t state_index = state_base + kd * dim + vd; + state[state_index] += normalized_k[kd] * delta; + output[vector_base + vd] += normalized_q[kd] * state[state_index]; + } + } + } + } + } +} + +std::vector makeValues(std::size_t count, float scale, float offset = 0.0F) { + std::vector values(count); + for (std::size_t i = 0; i < count; ++i) { + values[i] = offset + scale * static_cast((static_cast(i * 37U % 29U) - 14)); + } + return values; +} + +void ensureCpuContext() { + static const bool kContextInitialized = [] { + mllm::initializeContext(); + mllm::Context::instance().setCpuOpThreads(1); + return true; + }(); + (void)kContextInitialized; +} + +TEST(Ling3KDA, MatchesSafeGateReference) { + constexpr int kBatch = 2; + constexpr int kSequence = 5; + constexpr int kHeads = 3; + constexpr int kDim = 8; + const std::size_t vector_count = static_cast(kBatch) * kSequence * kHeads * kDim; + const std::size_t state_count = static_cast(kBatch) * kHeads * kDim * kDim; + + auto q = makeValues(vector_count, 0.017F); + auto k = makeValues(vector_count, -0.013F, 0.02F); + auto v = makeValues(vector_count, 0.011F, -0.03F); + auto gate = makeValues(vector_count, 0.019F, 0.1F); + auto beta = makeValues(static_cast(kBatch) * kSequence * kHeads, 0.01F, 0.45F); + auto a_log = makeValues(kHeads, 0.02F, 0.3F); + auto dt_bias = makeValues(static_cast(kHeads) * kDim, 0.015F, -0.1F); + auto expected_state = makeValues(state_count, 0.003F); + auto actual_state = expected_state; + std::vector expected_output(vector_count, 0.0F); + std::vector actual_output(vector_count, 0.0F); + + referenceKDA(q, k, v, gate, beta, a_log, dt_bias, expected_state, expected_output, kBatch, kSequence, kHeads, kDim, true, + -5.0F); + mllm::cpu::kda::kimiDeltaAttentionF32(q.data(), k.data(), v.data(), gate.data(), beta.data(), a_log.data(), dt_bias.data(), + actual_state.data(), actual_output.data(), kBatch, kSequence, kHeads, kDim, true, -5.0F, + 1); + + for (std::size_t i = 0; i < vector_count; ++i) { EXPECT_NEAR(actual_output[i], expected_output[i], 2.0e-6F); } + for (std::size_t i = 0; i < state_count; ++i) { EXPECT_NEAR(actual_state[i], expected_state[i], 2.0e-6F); } +} + +TEST(Ling3KDA, PrefillAndTokenwiseDecodeAreEquivalent) { + constexpr int kBatch = 1; + constexpr int kSequence = 7; + constexpr int kHeads = 2; + constexpr int kDim = 8; + const std::size_t vector_count = static_cast(kBatch) * kSequence * kHeads * kDim; + const std::size_t state_count = static_cast(kBatch) * kHeads * kDim * kDim; + auto q = makeValues(vector_count, 0.012F); + auto k = makeValues(vector_count, 0.009F, -0.02F); + auto v = makeValues(vector_count, -0.008F, 0.01F); + auto gate = makeValues(vector_count, 0.014F); + auto beta = makeValues(static_cast(kBatch) * kSequence * kHeads, 0.008F, 0.5F); + auto a_log = makeValues(kHeads, 0.03F, 0.2F); + auto dt_bias = makeValues(static_cast(kHeads) * kDim, 0.01F, -0.05F); + std::vector prefill_state(state_count, 0.0F); + std::vector decode_state(state_count, 0.0F); + std::vector prefill_output(vector_count, 0.0F); + std::vector decode_output(vector_count, 0.0F); + + mllm::cpu::kda::kimiDeltaAttentionF32(q.data(), k.data(), v.data(), gate.data(), beta.data(), a_log.data(), dt_bias.data(), + prefill_state.data(), prefill_output.data(), kBatch, kSequence, kHeads, kDim, true, + -5.0F, 1); + const std::size_t token_width = static_cast(kHeads) * kDim; + for (int token = 0; token < kSequence; ++token) { + const std::size_t offset = static_cast(token) * token_width; + mllm::cpu::kda::kimiDeltaAttentionF32(q.data() + offset, k.data() + offset, v.data() + offset, gate.data() + offset, + beta.data() + static_cast(token) * kHeads, a_log.data(), dt_bias.data(), + decode_state.data(), decode_output.data() + offset, kBatch, 1, kHeads, kDim, true, + -5.0F, 1); + } + + EXPECT_EQ(prefill_output, decode_output); + EXPECT_EQ(prefill_state, decode_state); +} + +TEST(Ling3KDA, RejectsInvalidSafeGateBound) { + std::vector values(4, 0.0F); + std::vector state(4, 0.0F); + EXPECT_THROW(mllm::cpu::kda::kimiDeltaAttentionF32(values.data(), values.data(), values.data(), values.data(), values.data(), + values.data(), values.data(), state.data(), values.data(), 1, 1, 1, 2, + true, 0.0F, 1), + std::invalid_argument); +} + +TEST(Ling3KDA, FunctionalDispatchReturnsOutputAndUpdatedState) { + ensureCpuContext(); + constexpr int kBatch = 1; + constexpr int kSequence = 3; + constexpr int kHeads = 2; + constexpr int kDim = 4; + const std::size_t vector_count = static_cast(kBatch) * kSequence * kHeads * kDim; + const std::size_t state_count = static_cast(kBatch) * kHeads * kDim * kDim; + + const auto q = makeValues(vector_count, 0.017F); + const auto k = makeValues(vector_count, -0.013F, 0.02F); + const auto v = makeValues(vector_count, 0.011F, -0.03F); + const auto gate = makeValues(vector_count, 0.019F, 0.1F); + const auto beta = makeValues(static_cast(kBatch) * kSequence * kHeads, 0.01F, 0.45F); + const auto a_log = makeValues(kHeads, 0.02F, 0.3F); + const auto dt_bias = makeValues(static_cast(kHeads) * kDim, 0.015F, -0.1F); + const auto initial_state = makeValues(state_count, 0.003F); + auto expected_state = initial_state; + std::vector expected_output(vector_count, 0.0F); + referenceKDA(q, k, v, gate, beta, a_log, dt_bias, expected_state, expected_output, kBatch, kSequence, kHeads, kDim, true, + -5.0F); + auto state_input = mllm::Tensor::fromVector(initial_state, {kBatch, kHeads, kDim, kDim}); + + auto [output, updated_state] = mllm::nn::functional::kimiDeltaAttention( + mllm::Tensor::fromVector(q, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(k, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(v, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(gate, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(beta, {kBatch, kSequence, kHeads}), mllm::Tensor::fromVector(a_log, {kHeads}), + mllm::Tensor::fromVector(dt_bias, {kHeads * kDim}), state_input, true, -5.0F); + + const auto actual_output = output.toVector(); + const auto actual_state = updated_state.toVector(); + EXPECT_EQ(state_input.toVector(), initial_state); + for (std::size_t i = 0; i < vector_count; ++i) { EXPECT_NEAR(actual_output[i], expected_output[i], 2.0e-6F); } + for (std::size_t i = 0; i < state_count; ++i) { EXPECT_NEAR(actual_state[i], expected_state[i], 2.0e-6F); } +} + +TEST(Ling3KDA, FunctionalDispatchMatchesKernelAtTinyProductionGeometry) { + ensureCpuContext(); + constexpr int kBatch = 1; + constexpr int kSequence = 49; + constexpr int kHeads = 16; + constexpr int kDim = 128; + const std::size_t vector_count = static_cast(kBatch) * kSequence * kHeads * kDim; + const std::size_t state_count = static_cast(kBatch) * kHeads * kDim * kDim; + const auto q = makeValues(vector_count, 0.0017F); + const auto k = makeValues(vector_count, -0.0013F, 0.002F); + const auto v = makeValues(vector_count, 0.0011F, -0.003F); + const auto gate = makeValues(vector_count, 0.0019F, 0.01F); + const auto beta = makeValues(static_cast(kBatch) * kSequence * kHeads, 0.001F, 0.45F); + const auto a_log = makeValues(kHeads, 0.002F, 0.3F); + const auto dt_bias = makeValues(static_cast(kHeads) * kDim, 0.0015F, -0.1F); + const auto initial_state = makeValues(state_count, 0.0003F); + auto expected_state = initial_state; + std::vector expected_output(vector_count, 0.0F); + mllm::cpu::kda::kimiDeltaAttentionF32(q.data(), k.data(), v.data(), gate.data(), beta.data(), a_log.data(), dt_bias.data(), + expected_state.data(), expected_output.data(), kBatch, kSequence, kHeads, kDim, true, + -5.0F, 8); + + auto [output, updated_state] = mllm::nn::functional::kimiDeltaAttention( + mllm::Tensor::fromVector(q, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(k, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(v, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(gate, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(beta, {kBatch, kSequence, kHeads}), mllm::Tensor::fromVector(a_log, {kHeads}), + mllm::Tensor::fromVector(dt_bias, {kHeads * kDim}), mllm::Tensor::fromVector(initial_state, {kBatch, kHeads, kDim, kDim}), + true, -5.0F); + + EXPECT_EQ(output.toVector(), expected_output); + EXPECT_EQ(updated_state.toVector(), expected_state); +} + +TEST(Ling3KDA, FunctionalInplaceStateAliasesAndMutatesInput) { + ensureCpuContext(); + constexpr int kBatch = 1; + constexpr int kSequence = 3; + constexpr int kHeads = 2; + constexpr int kDim = 4; + const std::size_t vector_count = static_cast(kBatch) * kSequence * kHeads * kDim; + const std::size_t state_count = static_cast(kBatch) * kHeads * kDim * kDim; + const auto q = makeValues(vector_count, 0.017F); + const auto k = makeValues(vector_count, -0.013F, 0.02F); + const auto v = makeValues(vector_count, 0.011F, -0.03F); + const auto gate = makeValues(vector_count, 0.019F, 0.1F); + const auto beta = makeValues(static_cast(kBatch) * kSequence * kHeads, 0.01F, 0.45F); + const auto a_log = makeValues(kHeads, 0.02F, 0.3F); + const auto dt_bias = makeValues(static_cast(kHeads) * kDim, 0.015F, -0.1F); + const auto initial_state = makeValues(state_count, 0.003F); + auto expected_state = initial_state; + std::vector expected_output(vector_count, 0.0F); + mllm::cpu::kda::kimiDeltaAttentionF32(q.data(), k.data(), v.data(), gate.data(), beta.data(), a_log.data(), dt_bias.data(), + expected_state.data(), expected_output.data(), kBatch, kSequence, kHeads, kDim, true, + -5.0F, 8); + auto state = mllm::Tensor::fromVector(initial_state, {kBatch, kHeads, kDim, kDim}); + + auto [output, updated_state] = mllm::nn::functional::kimiDeltaAttention( + mllm::Tensor::fromVector(q, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(k, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(v, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(gate, {kBatch, kSequence, kHeads, kDim}), + mllm::Tensor::fromVector(beta, {kBatch, kSequence, kHeads}), mllm::Tensor::fromVector(a_log, {kHeads}), + mllm::Tensor::fromVector(dt_bias, {kHeads * kDim}), state, true, -5.0F, true); + + EXPECT_EQ(updated_state.ptr(), state.ptr()); + EXPECT_EQ(output.toVector(), expected_output); + EXPECT_EQ(state.toVector(), expected_state); +} + +TEST(Ling3KDA, FunctionalDispatchProducesLinalgIR) { + ensureCpuContext(); + auto q = mllm::Tensor::empty({1, 2, 2, 4}, mllm::kFloat32, mllm::kCPU); + auto beta = mllm::Tensor::empty({1, 2, 2}, mllm::kFloat32, mllm::kCPU); + auto a_log = mllm::Tensor::empty({2}, mllm::kFloat32, mllm::kCPU); + auto dt_bias = mllm::Tensor::empty({8}, mllm::kFloat32, mllm::kCPU); + auto state = mllm::Tensor::empty({1, 2, 4, 4}, mllm::kFloat32, mllm::kCPU); + + mllm::ir::lowlevel::traceStart(); + auto outputs = mllm::nn::functional::kimiDeltaAttention(q, q, q, q, beta, a_log, dt_bias, state, true, -5.0F, true); + auto ir_context = mllm::ir::lowlevel::traceStop(); + + EXPECT_EQ(outputs[0].shape(), q.shape()); + EXPECT_EQ(outputs[1].shape(), state.shape()); + mllm::ir::linalg::LinalgIROp::ptr_t kda_ir; + for (const auto& op : ir_context->topLevelOp()->cast_()->getTopRegion()->ops()) { + if (op->isa_()) { + kda_ir = op->cast_(); + break; + } + } + ASSERT_NE(kda_ir, nullptr); + + const auto options = mllm::jit::binary::dumpLinalgIROptions(kda_ir); + EXPECT_EQ(options.at("safe_gate"), true); + EXPECT_FLOAT_EQ(options.at("lower_bound").get(), -5.0F); + EXPECT_EQ(options.at("state_inplace"), true); + const auto restored = + mllm::jit::interpreter::aopsFromJson(nlohmann::json{{"op_type", "KimiDeltaAttention"}, {"op_options", options}}); + const auto restored_kda = std::dynamic_pointer_cast(restored); + ASSERT_NE(restored_kda, nullptr); + EXPECT_EQ(restored_kda->options().safe_gate, true); + EXPECT_FLOAT_EQ(restored_kda->options().lower_bound, -5.0F); + EXPECT_EQ(restored_kda->options().state_inplace, true); +} + +} // namespace diff --git a/tests/cpu/Ling3RoPETest.cpp b/tests/cpu/Ling3RoPETest.cpp new file mode 100644 index 000000000..91182e453 --- /dev/null +++ b/tests/cpu/Ling3RoPETest.cpp @@ -0,0 +1,56 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/models/ling3/modeling_ling3.hpp" +#include "mllm/mllm.hpp" + +#include + +#include + +class Ling3RoPETest : public ::testing::Test { + protected: + static void SetUpTestSuite() { mllm::initializeContext(); } +}; + +TEST_F(Ling3RoPETest, AppliesOfficialAdjacentPairInterleave) { + auto input = mllm::Tensor::empty({1, 1, 2, 4}, mllm::kFloat32, mllm::kCPU).alloc(); + auto cos = mllm::Tensor::empty({1, 2, 4}, mllm::kFloat32, mllm::kCPU).alloc(); + auto sin = mllm::Tensor::empty({1, 2, 4}, mllm::kFloat32, mllm::kCPU).alloc(); + for (int index = 0; index < 8; ++index) { input.ptr()[index] = static_cast(index + 1); } + const float angles[] = {0.2F, -0.4F, 0.7F, 0.3F}; + for (int token = 0; token < 2; ++token) { + for (int pair = 0; pair < 2; ++pair) { + const float value = angles[token * 2 + pair]; + cos.ptr()[token * 4 + pair] = std::cos(value); + cos.ptr()[token * 4 + 2 + pair] = std::cos(value); + sin.ptr()[token * 4 + pair] = std::sin(value); + sin.ptr()[token * 4 + 2 + pair] = std::sin(value); + } + } + + const auto output = mllm::models::ling3::applyLing3InterleavedRoPE(input, cos, sin); + for (int token = 0; token < 2; ++token) { + for (int pair = 0; pair < 2; ++pair) { + const float even = input.ptr()[token * 4 + pair * 2]; + const float odd = input.ptr()[token * 4 + pair * 2 + 1]; + const float angle = angles[token * 2 + pair]; + EXPECT_NEAR(output.ptr()[token * 4 + pair], even * std::cos(angle) - odd * std::sin(angle), 1.0e-6F); + EXPECT_NEAR(output.ptr()[token * 4 + 2 + pair], odd * std::cos(angle) + even * std::sin(angle), 1.0e-6F); + } + } +} + +TEST_F(Ling3RoPETest, PadsMLAValuesWithoutChangingPayload) { + auto input = mllm::Tensor::empty({1, 2, 3, 4}, mllm::kFloat32, mllm::kCPU).alloc(); + for (int index = 0; index < 24; ++index) { input.ptr()[index] = static_cast(index + 1); } + const auto output = mllm::models::ling3::padLing3ValuesForCache(input, 6); + ASSERT_EQ(output.shape(), (mllm::Tensor::shape_t{1, 2, 3, 6})); + for (int vector = 0; vector < 6; ++vector) { + for (int dim = 0; dim < 4; ++dim) { + EXPECT_EQ(output.ptr()[vector * 6 + dim], input.ptr()[vector * 4 + dim]); + } + EXPECT_EQ(output.ptr()[vector * 6 + 4], 0.0F); + EXPECT_EQ(output.ptr()[vector * 6 + 5], 0.0F); + } +} diff --git a/tests/cpu/Ling3TokenizerTest.cpp b/tests/cpu/Ling3TokenizerTest.cpp new file mode 100644 index 000000000..5dc1ac5f1 --- /dev/null +++ b/tests/cpu/Ling3TokenizerTest.cpp @@ -0,0 +1,55 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/models/ling3/tokenization_ling3.hpp" +#include "mllm/mllm.hpp" + +#include + +#include +#include +#include + +namespace { + +std::string officialTokenizerPath() { + const char* value = std::getenv("LING3_OFFICIAL_TOKENIZER"); + return value == nullptr ? std::string() : std::string(value); +} + +std::vector ids(mllm::models::ling3::Ling3Tokenizer& tokenizer, const std::string& text) { + const auto tensor = tokenizer.convert2Ids(tokenizer.tokenize(text)); + return {tensor.ptr(), tensor.ptr() + tensor.numel()}; +} + +class Ling3TokenizerTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { mllm::initializeContext(); } +}; + +TEST_F(Ling3TokenizerTest, MatchesOfficialByteBPEAndNFCVectors) { + const auto path = officialTokenizerPath(); + if (path.empty()) { GTEST_SKIP() << "LING3_OFFICIAL_TOKENIZER is not set"; } + mllm::models::ling3::Ling3Tokenizer tokenizer(path); + EXPECT_EQ(ids(tokenizer, "Hello, world!"), (std::vector{14455, 11, 1931, 0})); + EXPECT_EQ(ids(tokenizer, "你好，世界！"), (std::vector{34355, 44291, 859})); + EXPECT_EQ(ids(tokenizer, "e\xCC\x81 café"), (std::vector{2900, 67656})); +} + +TEST_F(Ling3TokenizerTest, RendersOfficialSingleTurnThinkingTemplates) { + const auto path = officialTokenizerPath(); + if (path.empty()) { GTEST_SKIP() << "LING3_OFFICIAL_TOKENIZER is not set"; } + mllm::models::ling3::Ling3Tokenizer tokenizer(path); + auto enabled = tokenizer.convertMessage({"你好", "", true}).at("sequence"); + const std::vector enabled_expected = {157151, 90827, 157152, 14136, 5381, 6350, 366, 156895, 157151, 39, 116171, + 157152, 34355, 156895, 157151, 8469, 7342, 5468, 157152, 198, 156903}; + EXPECT_EQ(std::vector(enabled.ptr(), enabled.ptr() + enabled.numel()), enabled_expected); + + auto disabled = tokenizer.convertMessage({"Hello", "", false}).at("sequence"); + const std::vector disabled_expected = {157151, 90827, 157152, 14136, 5381, 6350, 928, 156895, + 157151, 39, 116171, 157152, 14455, 156895, 157151, 8469, + 7342, 5468, 157152, 198, 156903, 156904}; + EXPECT_EQ(std::vector(disabled.ptr(), disabled.ptr() + disabled.numel()), disabled_expected); +} + +} // namespace diff --git a/tests/models/CMakeLists.txt b/tests/models/CMakeLists.txt index c427145b8..e4ae8a325 100644 --- a/tests/models/CMakeLists.txt +++ b/tests/models/CMakeLists.txt @@ -1 +1,2 @@ add_subdirectory(qwen3_5) +add_subdirectory(ling3) diff --git a/tests/models/ling3/CMakeLists.txt b/tests/models/ling3/CMakeLists.txt new file mode 100644 index 000000000..e1c701546 --- /dev/null +++ b/tests/models/ling3/CMakeLists.txt @@ -0,0 +1,18 @@ +add_executable(Mllm-Test-Ling3-Config Ling3ConfigTest.cpp) +target_link_libraries(Mllm-Test-Ling3-Config PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-Ling3-Config PRIVATE ${MLLM_INCLUDE_DIR}) +target_compile_definitions(Mllm-Test-Ling3-Config + PRIVATE LING3_EXAMPLE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../../../examples/ling3") + +add_executable(Mllm-Test-Ling3-Tokenizer Ling3TokenizerTest.cpp) +target_link_libraries(Mllm-Test-Ling3-Tokenizer PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-Ling3-Tokenizer PRIVATE ${MLLM_INCLUDE_DIR}) + +add_executable(Mllm-Test-Ling3-RoPE Ling3RoPETest.cpp) +target_link_libraries(Mllm-Test-Ling3-RoPE PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-Ling3-RoPE PRIVATE ${MLLM_INCLUDE_DIR}) + +add_test(NAME Ling3ConfigFocused COMMAND Mllm-Test-Ling3-Config) +add_test(NAME Ling3TokenizerFocused COMMAND Mllm-Test-Ling3-Tokenizer) +add_test(NAME Ling3RoPEFocused COMMAND Mllm-Test-Ling3-RoPE) +set_tests_properties(Ling3ConfigFocused Ling3TokenizerFocused Ling3RoPEFocused PROPERTIES LABELS ling3) diff --git a/tests/models/ling3/Ling3ConfigTest.cpp b/tests/models/ling3/Ling3ConfigTest.cpp new file mode 100644 index 000000000..f90a770c0 --- /dev/null +++ b/tests/models/ling3/Ling3ConfigTest.cpp @@ -0,0 +1,30 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/models/ling3/configuration_ling3.hpp" + +#include + +#include +#include + +namespace { + +auto exampleDir() -> std::string { + const char* example_dir_override = std::getenv("MLLM_LING3_EXAMPLE_DIR"); + return example_dir_override == nullptr ? std::string(LING3_EXAMPLE_DIR) : std::string(example_dir_override); +} + +std::string ling3RuntimeConfigPath() { return exampleDir() + "/config_tiny_w4a32_kai.json"; } + +} // namespace + +TEST(Ling3Config, AcceptsPinnedTinyRuntimeContract) { + const mllm::models::ling3::Ling3Config config(ling3RuntimeConfigPath()); + EXPECT_TRUE(mllm::models::ling3::hasOfficialLing3TinyArchitecture(config)); + EXPECT_EQ(config.numFullAttentionLayers(), 6); + EXPECT_EQ(config.numKDALayers(), 18); + for (int layer = 0; layer < config.num_hidden_layers; ++layer) { + EXPECT_EQ(config.isFullAttentionLayer(layer), (layer + 1) % 4 == 0); + } +} diff --git a/tests/models/ling3/Ling3RoPETest.cpp b/tests/models/ling3/Ling3RoPETest.cpp new file mode 100644 index 000000000..91182e453 --- /dev/null +++ b/tests/models/ling3/Ling3RoPETest.cpp @@ -0,0 +1,56 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/models/ling3/modeling_ling3.hpp" +#include "mllm/mllm.hpp" + +#include + +#include + +class Ling3RoPETest : public ::testing::Test { + protected: + static void SetUpTestSuite() { mllm::initializeContext(); } +}; + +TEST_F(Ling3RoPETest, AppliesOfficialAdjacentPairInterleave) { + auto input = mllm::Tensor::empty({1, 1, 2, 4}, mllm::kFloat32, mllm::kCPU).alloc(); + auto cos = mllm::Tensor::empty({1, 2, 4}, mllm::kFloat32, mllm::kCPU).alloc(); + auto sin = mllm::Tensor::empty({1, 2, 4}, mllm::kFloat32, mllm::kCPU).alloc(); + for (int index = 0; index < 8; ++index) { input.ptr()[index] = static_cast(index + 1); } + const float angles[] = {0.2F, -0.4F, 0.7F, 0.3F}; + for (int token = 0; token < 2; ++token) { + for (int pair = 0; pair < 2; ++pair) { + const float value = angles[token * 2 + pair]; + cos.ptr()[token * 4 + pair] = std::cos(value); + cos.ptr()[token * 4 + 2 + pair] = std::cos(value); + sin.ptr()[token * 4 + pair] = std::sin(value); + sin.ptr()[token * 4 + 2 + pair] = std::sin(value); + } + } + + const auto output = mllm::models::ling3::applyLing3InterleavedRoPE(input, cos, sin); + for (int token = 0; token < 2; ++token) { + for (int pair = 0; pair < 2; ++pair) { + const float even = input.ptr()[token * 4 + pair * 2]; + const float odd = input.ptr()[token * 4 + pair * 2 + 1]; + const float angle = angles[token * 2 + pair]; + EXPECT_NEAR(output.ptr()[token * 4 + pair], even * std::cos(angle) - odd * std::sin(angle), 1.0e-6F); + EXPECT_NEAR(output.ptr()[token * 4 + 2 + pair], odd * std::cos(angle) + even * std::sin(angle), 1.0e-6F); + } + } +} + +TEST_F(Ling3RoPETest, PadsMLAValuesWithoutChangingPayload) { + auto input = mllm::Tensor::empty({1, 2, 3, 4}, mllm::kFloat32, mllm::kCPU).alloc(); + for (int index = 0; index < 24; ++index) { input.ptr()[index] = static_cast(index + 1); } + const auto output = mllm::models::ling3::padLing3ValuesForCache(input, 6); + ASSERT_EQ(output.shape(), (mllm::Tensor::shape_t{1, 2, 3, 6})); + for (int vector = 0; vector < 6; ++vector) { + for (int dim = 0; dim < 4; ++dim) { + EXPECT_EQ(output.ptr()[vector * 6 + dim], input.ptr()[vector * 4 + dim]); + } + EXPECT_EQ(output.ptr()[vector * 6 + 4], 0.0F); + EXPECT_EQ(output.ptr()[vector * 6 + 5], 0.0F); + } +} diff --git a/tests/models/ling3/Ling3TokenizerTest.cpp b/tests/models/ling3/Ling3TokenizerTest.cpp new file mode 100644 index 000000000..5dc1ac5f1 --- /dev/null +++ b/tests/models/ling3/Ling3TokenizerTest.cpp @@ -0,0 +1,55 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/models/ling3/tokenization_ling3.hpp" +#include "mllm/mllm.hpp" + +#include + +#include +#include +#include + +namespace { + +std::string officialTokenizerPath() { + const char* value = std::getenv("LING3_OFFICIAL_TOKENIZER"); + return value == nullptr ? std::string() : std::string(value); +} + +std::vector ids(mllm::models::ling3::Ling3Tokenizer& tokenizer, const std::string& text) { + const auto tensor = tokenizer.convert2Ids(tokenizer.tokenize(text)); + return {tensor.ptr(), tensor.ptr() + tensor.numel()}; +} + +class Ling3TokenizerTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { mllm::initializeContext(); } +}; + +TEST_F(Ling3TokenizerTest, MatchesOfficialByteBPEAndNFCVectors) { + const auto path = officialTokenizerPath(); + if (path.empty()) { GTEST_SKIP() << "LING3_OFFICIAL_TOKENIZER is not set"; } + mllm::models::ling3::Ling3Tokenizer tokenizer(path); + EXPECT_EQ(ids(tokenizer, "Hello, world!"), (std::vector{14455, 11, 1931, 0})); + EXPECT_EQ(ids(tokenizer, "你好，世界！"), (std::vector{34355, 44291, 859})); + EXPECT_EQ(ids(tokenizer, "e\xCC\x81 café"), (std::vector{2900, 67656})); +} + +TEST_F(Ling3TokenizerTest, RendersOfficialSingleTurnThinkingTemplates) { + const auto path = officialTokenizerPath(); + if (path.empty()) { GTEST_SKIP() << "LING3_OFFICIAL_TOKENIZER is not set"; } + mllm::models::ling3::Ling3Tokenizer tokenizer(path); + auto enabled = tokenizer.convertMessage({"你好", "", true}).at("sequence"); + const std::vector enabled_expected = {157151, 90827, 157152, 14136, 5381, 6350, 366, 156895, 157151, 39, 116171, + 157152, 34355, 156895, 157151, 8469, 7342, 5468, 157152, 198, 156903}; + EXPECT_EQ(std::vector(enabled.ptr(), enabled.ptr() + enabled.numel()), enabled_expected); + + auto disabled = tokenizer.convertMessage({"Hello", "", false}).at("sequence"); + const std::vector disabled_expected = {157151, 90827, 157152, 14136, 5381, 6350, 928, 156895, + 157151, 39, 116171, 157152, 14455, 156895, 157151, 8469, + 7342, 5468, 157152, 198, 156903, 156904}; + EXPECT_EQ(std::vector(disabled.ptr(), disabled.ptr() + disabled.numel()), disabled_expected); +} + +} // namespace diff --git a/tests/nn/CMakeLists.txt b/tests/nn/CMakeLists.txt index bd003caac..b9a96593f 100644 --- a/tests/nn/CMakeLists.txt +++ b/tests/nn/CMakeLists.txt @@ -26,5 +26,10 @@ add_executable(Mllm-Test-Nn-GatedDeltaRule GatedDeltaRuleTest.cpp) target_link_libraries(Mllm-Test-Nn-GatedDeltaRule PRIVATE gtest_main MllmRT MllmCPUBackend) target_include_directories(Mllm-Test-Nn-GatedDeltaRule PRIVATE ${MLLM_INCLUDE_DIR}) +add_executable(Mllm-Test-Nn-KimiDeltaAttention KimiDeltaAttentionTest.cpp) +target_link_libraries(Mllm-Test-Nn-KimiDeltaAttention PRIVATE gtest_main MllmRT MllmCPUBackend) +target_include_directories(Mllm-Test-Nn-KimiDeltaAttention PRIVATE ${MLLM_INCLUDE_DIR}) + include(GoogleTest) add_test(NAME GatedDeltaRuleFocused COMMAND Mllm-Test-Nn-GatedDeltaRule) +add_test(NAME KimiDeltaAttentionFocused COMMAND Mllm-Test-Nn-KimiDeltaAttention) diff --git a/tests/nn/KimiDeltaAttentionTest.cpp b/tests/nn/KimiDeltaAttentionTest.cpp new file mode 100644 index 000000000..3e5614e06 --- /dev/null +++ b/tests/nn/KimiDeltaAttentionTest.cpp @@ -0,0 +1,303 @@ +// 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 KimiDeltaAttentionTest : public testing::Test { + protected: + static void SetUpTestSuite() { mllm::initializeContext(); } +}; + +class KimiDeltaAttentionModule final : public mllm::nn::Module { + public: + KimiDeltaAttentionModule(std::string name, bool safe_gate, float lower_bound, bool state_inplace) : Module(std::move(name)) { + kda_ = reg("kda", safe_gate, lower_bound, state_inplace); + } + + std::vector forward(const std::vector& inputs, const std::vector&) override { + auto [output, state] = kda_(inputs[0], inputs[1], inputs[2], inputs[3], inputs[4], inputs[5], inputs[6], inputs[7]); + return {output, state}; + } + + private: + mllm::nn::KimiDeltaAttention kda_; +}; + +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 patterned(const Tensor::shape_t& shape, float scale, float offset = 0.0F) { + auto tensor = Tensor::empty(shape, mllm::kFloat32, mllm::kCPU).alloc(); + for (int index = 0; index < tensor.numel(); ++index) { + tensor.ptr()[index] = std::sin(static_cast(index + 1) * scale) + offset; + } + return tensor; +} + +float stableSigmoid(float value) { + if (value >= 0.0F) { + const float exp_value = std::exp(-value); + return 1.0F / (1.0F + exp_value); + } + const float exp_value = std::exp(value); + return exp_value / (1.0F + exp_value); +} + +// Independent scalar reference of the public contract: +// q/k/v/gate [B, S, H, D], beta [B, S, H] (post-sigmoid), a_log [H], +// dt_bias [H * D], state [B, H, D, D] -> output [B, S, H, D], new state. +std::pair, std::vector> referenceKimiDeltaAttention( + const std::vector& q, const std::vector& k, const std::vector& v, const std::vector& gate, + const std::vector& beta, const std::vector& a_log, const std::vector& dt_bias, + std::vector state, int batch, int sequence, int heads, int dim, bool safe_gate, float lower_bound) { + const float query_dim_scale = 1.0F / std::sqrt(static_cast(dim)); + std::vector output(static_cast(batch) * sequence * heads * dim, 0.0F); + std::vector normalized_q(dim); + std::vector normalized_k(dim); + std::vector prediction(dim); + for (int b = 0; b < batch; ++b) { + for (int h = 0; h < heads; ++h) { + const std::size_t state_base = (static_cast(b) * heads + h) * dim * dim; + const float a_scale = std::exp(a_log[h]); + for (int s = 0; s < sequence; ++s) { + const std::size_t vector_base = ((static_cast(b) * sequence + s) * heads + h) * dim; + float q_norm_sq = 0.0F; + float k_norm_sq = 0.0F; + for (int d = 0; d < dim; ++d) { + q_norm_sq += q[vector_base + d] * q[vector_base + d]; + k_norm_sq += k[vector_base + d] * k[vector_base + d]; + } + const float q_scale = query_dim_scale / std::sqrt(q_norm_sq + 1.0e-6F); + const float k_scale = 1.0F / std::sqrt(k_norm_sq + 1.0e-6F); + for (int d = 0; d < dim; ++d) { + normalized_q[d] = q[vector_base + d] * q_scale; + normalized_k[d] = k[vector_base + d] * k_scale; + prediction[d] = 0.0F; + } + for (int kd = 0; kd < dim; ++kd) { + const float gate_input = gate[vector_base + kd] + dt_bias[h * dim + kd]; + const float log_decay = safe_gate ? lower_bound * stableSigmoid(a_scale * gate_input) + : -a_scale * std::log1p(std::exp(gate_input)); + const float decay = std::exp(log_decay); + for (int vd = 0; vd < dim; ++vd) { + const std::size_t state_index = state_base + kd * dim + vd; + state[state_index] *= decay; + prediction[vd] += normalized_k[kd] * state[state_index]; + } + } + const float beta_value = beta[(static_cast(b) * sequence + s) * heads + h]; + for (int vd = 0; vd < dim; ++vd) { + const float delta = beta_value * (v[vector_base + vd] - prediction[vd]); + for (int kd = 0; kd < dim; ++kd) { + const std::size_t state_index = state_base + kd * dim + vd; + state[state_index] += normalized_k[kd] * delta; + output[vector_base + vd] += normalized_q[kd] * state[state_index]; + } + } + } + } + } + return {output, state}; +} + +void expectNear(const Tensor& actual, const std::vector& expected, float tolerance) { + ASSERT_EQ(actual.numel(), expected.size()); + for (int index = 0; index < actual.numel(); ++index) { + ASSERT_NEAR(actual.ptr()[index], expected[index], tolerance) << "index " << index; + } +} + +struct Case { + int batch; + int sequence; + int heads; + int dim; + bool safe_gate; + float lower_bound; + float tolerance; +}; + +void runEagerReferenceCase(const Case& test_case, const std::string& module_name) { + SCOPED_TRACE(::testing::Message() << module_name << " B=" << test_case.batch << " S=" << test_case.sequence + << " H=" << test_case.heads << " D=" << test_case.dim + << " safe_gate=" << test_case.safe_gate); + const Tensor::shape_t vector_shape = {test_case.batch, test_case.sequence, test_case.heads, test_case.dim}; + auto q = patterned(vector_shape, 0.03F); + auto k = patterned(vector_shape, 0.05F, 0.1F); + auto v = patterned(vector_shape, 0.07F, -0.2F); + auto gate = patterned(vector_shape, 0.09F, -0.2F); + auto beta = patterned({test_case.batch, test_case.sequence, test_case.heads}, 0.11F, 0.5F); + auto a_log = patterned({test_case.heads}, 0.13F, -0.4F); + auto dt_bias = patterned({test_case.heads * test_case.dim}, 0.15F, -0.1F); + auto state = patterned({test_case.batch, test_case.heads, test_case.dim, test_case.dim}, 0.017F); + const auto state_before = state.toVector(); + const auto [expected_output, expected_state] = referenceKimiDeltaAttention( + q.toVector(), k.toVector(), v.toVector(), gate.toVector(), beta.toVector(), + a_log.toVector(), dt_bias.toVector(), state_before, test_case.batch, test_case.sequence, test_case.heads, + test_case.dim, test_case.safe_gate, test_case.lower_bound); + + KimiDeltaAttentionModule module(module_name, test_case.safe_gate, test_case.lower_bound, /*state_inplace=*/false); + const auto outputs = module(q, k, v, gate, beta, a_log, dt_bias, state); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].shape(), vector_shape); + EXPECT_EQ(outputs[1].shape(), state.shape()); + EXPECT_NE(outputs[1].ptr(), state.ptr()); + EXPECT_EQ(state.toVector(), state_before); + expectNear(outputs[0], expected_output, test_case.tolerance); + expectNear(outputs[1], expected_state, test_case.tolerance); +} + +TEST_F(KimiDeltaAttentionTest, EagerMatchesIndependentReferenceAndPreservesInputState) { + runEagerReferenceCase({2, 5, 3, 8, true, -5.0F, 2.0e-6F}, "kimi_delta_attention_safe_gate"); + runEagerReferenceCase({1, 4, 2, 8, false, 0.0F, 2.0e-6F}, "kimi_delta_attention_softplus_gate"); + runEagerReferenceCase({1, 3, 2, 6, true, -2.5F, 2.0e-6F}, "kimi_delta_attention_custom_bound"); +} + +TEST_F(KimiDeltaAttentionTest, EagerMatchesIndependentReferenceAtProductionHeadGeometry) { + // Ling-3.0-tiny geometry: 16 heads of 128 lanes; the batched lane helpers + // and multi-threaded (batch, head) partition are all exercised here. + runEagerReferenceCase({1, 49, 16, 128, true, -5.0F, 5.0e-5F}, "kimi_delta_attention_production"); +} + +TEST_F(KimiDeltaAttentionTest, InplaceStateOutputAliasesAndMutatesInput) { + auto q = patterned({1, 3, 2, 4}, 0.03F); + auto k = patterned(q.shape(), 0.05F); + auto v = patterned(q.shape(), 0.07F); + auto gate = patterned(q.shape(), 0.09F, -0.2F); + auto beta = patterned({1, 3, 2}, 0.11F, 0.5F); + auto a_log = patterned({2}, 0.13F, -0.4F); + auto dt_bias = patterned({8}, 0.15F, -0.1F); + auto state = patterned({1, 2, 4, 4}, 0.017F); + const auto state_before = state.toVector(); + const auto* state_storage = state.ptr(); + + KimiDeltaAttentionModule module("kimi_delta_attention_inplace", true, -5.0F, /*state_inplace=*/true); + const auto outputs = module(q, k, v, gate, beta, a_log, dt_bias, state); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[1].ptr(), state_storage); + EXPECT_NE(state.toVector(), state_before); + EXPECT_EQ(outputs[1].toVector(), state.toVector()); +} + +TEST_F(KimiDeltaAttentionTest, ChunkedPrefillAndDecodeMatchOneShot) { + constexpr int kSequence = 7; + constexpr int kHeads = 2; + constexpr int kDim = 8; + auto q = patterned({1, kSequence, kHeads, kDim}, 0.03F); + auto k = patterned(q.shape(), 0.05F, 0.1F); + auto v = patterned(q.shape(), 0.07F, -0.2F); + auto gate = patterned(q.shape(), 0.09F, -0.2F); + auto beta = patterned({1, kSequence, kHeads}, 0.11F, 0.5F); + auto a_log = patterned({kHeads}, 0.13F, -0.4F); + auto dt_bias = patterned({kHeads * kDim}, 0.15F, -0.1F); + const auto initial_state = patterned({1, kHeads, kDim, kDim}, 0.017F).toVector(); + + KimiDeltaAttentionModule one_shot("kimi_delta_attention_one_shot", true, -5.0F, false); + const auto expected = + one_shot(q, k, v, gate, beta, a_log, dt_bias, Tensor::fromVector(initial_state, {1, kHeads, kDim, kDim})); + const auto expected_output = expected[0].toVector(); + + KimiDeltaAttentionModule chunked("kimi_delta_attention_chunked", true, -5.0F, false); + auto state = Tensor::fromVector(initial_state, {1, kHeads, kDim, kDim}); + std::vector chunked_output; + const std::vector> chunks = {{0, 3}, {3, 4}, {4, 5}, {5, 7}}; + for (const auto& [begin, end] : chunks) { + const auto length = end - begin; + // Token slices [begin, end) of a [1, S, H, D] activation or a [1, S, H] gate. + const auto slice = [&](const Tensor& tensor, int width, const Tensor::shape_t& shape) { + const auto values = tensor.toVector(); + const auto first = values.begin() + static_cast(begin) * width; + return Tensor::fromVector(std::vector(first, first + static_cast(length) * width), shape); + }; + const Tensor::shape_t vector_shape = {1, length, kHeads, kDim}; + const auto outputs = chunked(slice(q, kHeads * kDim, vector_shape), slice(k, kHeads * kDim, vector_shape), + slice(v, kHeads * kDim, vector_shape), slice(gate, kHeads * kDim, vector_shape), + slice(beta, kHeads, {1, length, kHeads}), a_log, dt_bias, state); + const auto values = outputs[0].toVector(); + chunked_output.insert(chunked_output.end(), values.begin(), values.end()); + state = outputs[1]; + } + + EXPECT_EQ(chunked_output, expected_output); + EXPECT_EQ(state.toVector(), expected[1].toVector()); +} + +TEST_F(KimiDeltaAttentionTest, RejectsInvalidGeometryAndOptions) { + auto q = patterned({1, 2, 2, 4}, 0.03F); + auto beta = patterned({1, 2, 2}, 0.11F, 0.5F); + auto a_log = patterned({2}, 0.13F, -0.4F); + auto dt_bias = patterned({8}, 0.15F, -0.1F); + auto state = patterned({1, 2, 4, 4}, 0.017F); + + KimiDeltaAttentionModule module("kimi_delta_attention_invalid", true, -5.0F, false); + // beta must be [B, S, H]. + EXPECT_THROW((void)module(q, q, q, q, patterned({1, 2, 2, 1}, 0.1F), a_log, dt_bias, state), std::invalid_argument); + // state must be [B, H, D, D]. + EXPECT_THROW((void)module(q, q, q, q, beta, a_log, dt_bias, patterned({1, 2, 4, 3}, 0.1F)), std::invalid_argument); + // dt_bias must hold H * D values. + EXPECT_THROW((void)module(q, q, q, q, beta, a_log, patterned({2}, 0.1F), state), std::invalid_argument); + // The safe gate requires a finite, strictly negative lower bound. + KimiDeltaAttentionModule invalid_bound("kimi_delta_attention_invalid_bound", true, 0.0F, false); + EXPECT_THROW((void)invalid_bound(q, q, q, q, beta, a_log, dt_bias, state), std::invalid_argument); +} + +TEST_F(KimiDeltaAttentionTest, TraceAndSerializationPreserveGateAndStateOptions) { + KimiDeltaAttentionModule module("kimi_delta_attention_trace", true, -3.5F, true); + auto ir_context = mllm::ir::trace( + module, Tensor::empty({1, 2, 2, 4}, mllm::kFloat32, mllm::kCPU), Tensor::empty({1, 2, 2, 4}, mllm::kFloat32, mllm::kCPU), + Tensor::empty({1, 2, 2, 4}, mllm::kFloat32, mllm::kCPU), Tensor::empty({1, 2, 2, 4}, mllm::kFloat32, mllm::kCPU), + Tensor::empty({1, 2, 2}, mllm::kFloat32, mllm::kCPU), Tensor::empty({2}, mllm::kFloat32, mllm::kCPU), + Tensor::empty({8}, mllm::kFloat32, mllm::kCPU), Tensor::empty({1, 2, 4, 4}, mllm::kFloat32, mllm::kCPU)); + auto op = findOp(ir_context->topLevelOp()); + ASSERT_NE(op, nullptr); + EXPECT_EQ(op->getAOp()->getOpType(), mllm::OpTypes::kKimiDeltaAttention); + EXPECT_EQ(op->inputs().size(), 8); + EXPECT_EQ(op->outputs().size(), 2); + + const auto options = mllm::jit::binary::dumpLinalgIROptions(op); + EXPECT_TRUE(options.at("safe_gate").get()); + EXPECT_FLOAT_EQ(options.at("lower_bound").get(), -3.5F); + EXPECT_TRUE(options.at("state_inplace").get()); + + const auto restored = mllm::jit::interpreter::aopsFromJson( + nlohmann::json{{"op_type", "KimiDeltaAttention"}, {"backend", "CPU"}, {"op_options", options}}); + ASSERT_NE(restored, nullptr); + EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kKimiDeltaAttention); + const auto restored_options = std::static_pointer_cast(restored)->options(); + EXPECT_TRUE(restored_options.safe_gate); + EXPECT_FLOAT_EQ(restored_options.lower_bound, -3.5F); + EXPECT_TRUE(restored_options.state_inplace); +} + +} // namespace