From 4cdbd107d1709fbbd2dea4a022a9d601cce254a9 Mon Sep 17 00:00:00 2001 From: xianyue <2869219499@qq.com> Date: Sat, 19 Sep 2026 09:03:47 +0800 Subject: [PATCH 1/2] feat(model): support LFM2-1.2B --- LFM2_REPORT.md | 183 +++++++++++ csrc/cache/kv_cache.cpp | 5 +- .../layers/causal_lm_templates/text_model.hpp | 6 +- .../layers/quantization/none_quantization.cpp | 11 +- csrc/models/lfm2/lfm2_allocate_cache.cpp | 101 ++++++ csrc/models/lfm2/lfm2_allocate_cache.hpp | 24 ++ csrc/models/lfm2/lfm2_decoder_layer.cpp | 72 +++++ csrc/models/lfm2/lfm2_decoder_layer.hpp | 44 +++ csrc/models/lfm2/lfm2_for_causal_lm.cpp | 108 +++++++ csrc/models/lfm2/lfm2_for_causal_lm.hpp | 28 ++ csrc/models/lfm2/lfm2_mlp.hpp | 27 ++ csrc/models/lfm2/lfm2_rms_norm.hpp | 91 ++++++ csrc/models/lfm2/lfm2_short_conv.cpp | 221 +++++++++++++ csrc/models/lfm2/lfm2_short_conv.hpp | 49 +++ python/infinilm/infer_engine.py | 3 + python/infinilm/modeling_utils.py | 26 ++ python/infinilm/processors/lfm2_processor.py | 76 +++++ test/models/lfm2/native_lfm2_tiny_smoke.py | 182 +++++++++++ test/models/lfm2/reference_lfm2.py | 222 +++++++++++++ test/models/lfm2/reference_lfm2_real.py | 129 ++++++++ test/models/lfm2/run_infinilm_lfm2_real.py | 295 ++++++++++++++++++ .../lfm2/test_low_precision_contract.py | 61 ++++ test/models/lfm2/test_run_config.py | 38 +++ test/models/lfm2/test_short_conv_contract.py | 125 ++++++++ test/models/lfm2/test_state_routing.py | 40 +++ test/models/lfm2/test_weight_remap.py | 93 ++++++ test/models/lfm2/tiny_config.json | 35 +++ .../lfm2/validate_lfm2_cache_and_timing.py | 264 ++++++++++++++++ 28 files changed, 2549 insertions(+), 10 deletions(-) create mode 100644 LFM2_REPORT.md create mode 100644 csrc/models/lfm2/lfm2_allocate_cache.cpp create mode 100644 csrc/models/lfm2/lfm2_allocate_cache.hpp create mode 100644 csrc/models/lfm2/lfm2_decoder_layer.cpp create mode 100644 csrc/models/lfm2/lfm2_decoder_layer.hpp create mode 100644 csrc/models/lfm2/lfm2_for_causal_lm.cpp create mode 100644 csrc/models/lfm2/lfm2_for_causal_lm.hpp create mode 100644 csrc/models/lfm2/lfm2_mlp.hpp create mode 100644 csrc/models/lfm2/lfm2_rms_norm.hpp create mode 100644 csrc/models/lfm2/lfm2_short_conv.cpp create mode 100644 csrc/models/lfm2/lfm2_short_conv.hpp create mode 100644 python/infinilm/processors/lfm2_processor.py create mode 100644 test/models/lfm2/native_lfm2_tiny_smoke.py create mode 100644 test/models/lfm2/reference_lfm2.py create mode 100644 test/models/lfm2/reference_lfm2_real.py create mode 100644 test/models/lfm2/run_infinilm_lfm2_real.py create mode 100644 test/models/lfm2/test_low_precision_contract.py create mode 100644 test/models/lfm2/test_run_config.py create mode 100644 test/models/lfm2/test_short_conv_contract.py create mode 100644 test/models/lfm2/test_state_routing.py create mode 100644 test/models/lfm2/test_weight_remap.py create mode 100644 test/models/lfm2/tiny_config.json create mode 100644 test/models/lfm2/validate_lfm2_cache_and_timing.py diff --git a/LFM2_REPORT.md b/LFM2_REPORT.md new file mode 100644 index 000000000..4be464bb4 --- /dev/null +++ b/LFM2_REPORT.md @@ -0,0 +1,183 @@ +# InfiniLM LFM2-1.2B 适配报告 + +## 1. 项目内容 + +本项目在 InfiniLM 中新增 `LiquidAI/LFM2-1.2B` 支持。LFM2 不是普通的全 Attention 模型，其 16 个 Decoder 层由 10 个 ShortConv 层和 6 个全 Attention 层交错组成，因此除了 Attention KV Cache，还需要维护 ShortConv 的卷积状态。 + +本次实现包括： + +- 注册 `model_type=lfm2`，将官方配置转换为 InfiniLM 使用的模型配置。 +- 实现 LFM2 Decoder、RMSNorm、SwiGLU MLP 和长度为 3 的 gated depthwise ShortConv。 +- 复用 Qwen3 GQA Attention，并按照官方 `layer_types` 组装混合 Decoder。 +- 为 Static/Paged 两种缓存模式分配并路由 Attention KV Cache 和 ShortConv State Cache。 +- 增加官方 Safetensors 权重名称到 InfiniLM 参数树的映射。 +- 增加真实模型、同权重 tiny 模型、缓存/reset、低精度计算和性能记录脚本。 + +## 2. 实现思路 + +### 2.1 混合 Decoder + +配置加载阶段根据 `full_attn_idxs` 生成每层的 `layer_types`。`full_attention` 层复用现有 Qwen3 Attention，`short_conv` 层使用新增的 `Lfm2ShortConv`。两类层共用 LFM2 RMSNorm 和 SwiGLU MLP，从而只新增 LFM2 特有结构，尽量复用现有基础设施。 + +### 2.2 ShortConv + +ShortConv 首先通过 `in_proj` 生成三个分支 `B`、`C` 和 `x`，计算： + +```text +y = out_proj(C * depthwise_causal_conv1d(B * x)) +``` + +Prefill 阶段使用滑动窗口和 batched Matmul 实现长度为 3 的深度卷积；单 Token Decode 阶段读取每个请求对应的历史状态，只计算当前 Token，并把最后两个时间步写回 Conv State Cache。 + +低精度路径显式区分 Prefill 和 Decode 的舍入边界，以复现 Transformers 参考实现的 BF16 计算顺序。 + +### 2.3 双缓存与请求隔离 + +Attention 层继续使用现有 KV Cache。ShortConv 层单独分配 `[state_pool, hidden_size, kernel_size - 1]` 状态张量，通过请求的初始/最终状态索引读取和写回。Static Cache 预留零历史行，确保新请求不会读取上一个请求的卷积状态；Paged Cache 则按请求索引保存状态。 + +### 2.4 权重映射 + +官方 LFM2 的 Attention Norm、输出投影、FFN 和最终 Norm 名称与 InfiniLM 参数树不完全相同。`_remap_lfm2` 在加载时完成名称转换，同时保留 ShortConv 原有权重名称。真实模型权重映射已经完成闭环验证。 + +## 3. 复现流程 + +### 3.1 环境 + +主要 NVIDIA 验证环境： + +```text +GPU: NVIDIA RTX 4090 24 GB +CUDA Toolkit: 12.8 +Model: LiquidAI/LFM2-1.2B +Dtype: BF16 +Decoding: greedy argmax +``` + +先按 InfiniCore README 编译并安装 NVIDIA 后端，并设置： + +```bash +export CUDA_HOME=/usr/local/cuda-12.8 +export INFINI_ROOT=/data/InfiniTensor/install/nvidia +export LD_LIBRARY_PATH="$INFINI_ROOT/lib:$CUDA_HOME/lib64:$LD_LIBRARY_PATH" +``` + +然后构建并安装 InfiniLM： + +```bash +xmake f -y -c -m release +xmake build -j4 _infinilm +xmake install _infinilm +python -m pip install --no-build-isolation --no-deps -e . +``` + +### 3.2 纯 Python 合同测试 + +```bash +PYTHONPATH=test/models/lfm2:python python -m unittest \ + test.models.lfm2.test_weight_remap \ + test.models.lfm2.test_state_routing \ + test.models.lfm2.test_short_conv_contract \ + test.models.lfm2.test_low_precision_contract \ + test.models.lfm2.test_run_config +``` + +当前结果：`16/16` 通过。 + +### 3.3 生成 Transformers 参考 + +分别对三个 Prompt 执行： + +```bash +python test/models/lfm2/reference_lfm2_real.py \ + --model /path/to/LFM2-1.2B \ + --prompt "Who are you?" \ + --max-new-tokens 16 \ + --device cuda \ + --output artifacts/lfm2_transformers_cuda.json +``` + +另外两组 Prompt 为： + +```text +请用一句中文介绍你自己。 +Explain in three short points why recurrent state can reduce decoding work. +``` + +### 3.4 InfiniLM Static/Paged 验证 + +Static Cache： + +```bash +python test/models/lfm2/run_infinilm_lfm2_real.py \ + --model /path/to/LFM2-1.2B \ + --device cuda \ + --cache-type static \ + --max-new-tokens 16 \ + --repeat 2 \ + --reference artifacts/lfm2_transformers_cuda.json \ + --extra-reference artifacts/lfm2_transformers_cuda_zh.json \ + --extra-reference artifacts/lfm2_transformers_cuda_long.json \ + --output artifacts/lfm2_static_gate.json +``` + +Paged Cache 使用相同命令，将 `--cache-type` 改为 `paged`。 + +## 4. 复现结果 + +### 4.1 NVIDIA RTX 4090 + +| 验收项 | Static | Paged | +|---|---:|---:| +| 三 Prompt Transformers 16-token 精确对齐 | 3/3 通过 | 3/3 通过 | +| A/B/C/A/B/C 重复请求 | 通过 | 通过 | +| 重复结果一致 | 通过 | 通过 | +| 208-token 输入、64 步固定长度压力测试 | 通过 | 通过 | +| Paged BF16 cache/full argmax 对照 | 不适用 | 56/56 | +| Static BF16 cache/full argmax 对照 | 54/56，未完全通过 | 不适用 | + +Prompt `Who are you?` 的 16 个生成 Token 与 Transformers 一致： + +```text +[550, 1283, 902, 14009, 6544, 16701, 5237, 811, + 1801, 7039, 916, 768, 6266, 3795, 803, 10003] +``` + +一次 RTX 4090 BF16 基线记录如下。计时范围只包括 engine forward 和设备同步，不包含模型加载、Tokenizer、元数据构造和输出回传，因此不是服务端端到端 TTFT： + +| 缓存模式 | 208-token Prefill | 单 Token Decode 均值 | Decode tokens/s | 设备显存采样 | +|---|---:|---:|---:|---:| +| Static | 7.43 ms | 6.04 ms | 165.49 | 3096 MiB | +| Paged | 6.63 ms | 4.50 ms | 222.25 | 3096 MiB | + +这些数据是单机基线，不是优化前后对比。 + +### 4.2 CPU + +CPU 原生构建、LFM2 模型工厂、参数树和 tiny F32 Full/Prefill/Decode 已通过。当前最终源码尚未重新执行真实 1.2B CPU 全量回归，因此 CPU 不标记为真实模型完整支持。 + +### 4.3 昇腾 910B1 + +当前为部分支持： + +- Ascend Runtime、设备复制和 LFM2 所需基础算子已完成验证。 +- tiny LFM2 F32 Full 与 Prefill+Decode 一致。 +- 真实 LFM2-1.2B 可以创建缓存、加载权重并进入第 0 层。 +- 真实 BF16 在 ShortConv 的 Linear/GEMM 阶段仍会触发 CANN 同步错误，尚未获得可信 logits 和 Token 对齐结果。 + +因此本报告不把昇腾标记为端到端支持。相关 InfiniCore Ascend 实验代码也尚未达到可合并状态。 + +## 5. 已知限制 + +- 当前正确性主线限定 `batch_size=1`；多个不同长度请求的 packed Prefill 尚未实现 ShortConv 分段状态更新。 +- tiny BF16 严格 logits 阈值仍未完全通过，虽然真实模型三 Prompt 的 greedy Token 已对齐。 +- Static BF16 的 cache/full 对照为 54/56 argmax 一致，仍有两个低决策间隔步骤发生分叉；Paged 路径为 56/56。 +- 尚未执行服务端压力测试、MMLU/C-Eval 和完整 CI。 +- 昇腾真实 BF16 端到端推理尚未完成。 + +## 6. 平台状态结论 + +| 平台 | 状态 | +|---|---| +| NVIDIA RTX 4090 / CUDA 12.8 | 主要功能通过；Static/Paged 三 Prompt 生成与 Transformers 精确对齐 | +| CPU | 构建、模型创建和 tiny F32 通过；真实 1.2B 最终回归未执行 | +| Ascend 910B1 / CANN 9.0 | Runtime、基础算子、tiny F32 和模型加载通过；真实 BF16 Linear/GEMM 未完成 | diff --git a/csrc/cache/kv_cache.cpp b/csrc/cache/kv_cache.cpp index 2a7780090..60a9e0329 100644 --- a/csrc/cache/kv_cache.cpp +++ b/csrc/cache/kv_cache.cpp @@ -58,7 +58,7 @@ infinicore::Tensor create_layer_kv_cache( size_t cache_len = (config.max_cache_len() == std::numeric_limits::max() || config.max_cache_len() == 0 ? max_positional_embedding : config.max_cache_len()); // Allocate KV cache - infinicore::Tensor kv_cache = infinicore::Tensor::empty( + infinicore::Tensor kv_cache = infinicore::Tensor::zeros( {2, rank_batch_size, num_rank_k_heads, @@ -66,9 +66,6 @@ infinicore::Tensor create_layer_kv_cache( kv_dim}, dtype, rank_info.device); - set_zeros(kv_cache); - - infinicore::context::syncStream(); return kv_cache; } diff --git a/csrc/layers/causal_lm_templates/text_model.hpp b/csrc/layers/causal_lm_templates/text_model.hpp index 49b60d7f4..cfed9b89f 100644 --- a/csrc/layers/causal_lm_templates/text_model.hpp +++ b/csrc/layers/causal_lm_templates/text_model.hpp @@ -24,7 +24,7 @@ namespace infinilm::layers::causal_lm_templates { * * @tparam DecoderLayer The decoder layer type (e.g., Qwen3DecoderLayer) */ -template +template class TextModel : public infinicore::nn::Module { public: TextModel(std::shared_ptr model_config, @@ -56,7 +56,7 @@ class TextModel : public infinicore::nn::Module { } if (is_last_pp_stage()) { - norm_ = this->register_module("norm", hidden_size_, rms_norm_eps, dtype, device); + norm_ = this->register_module("norm", hidden_size_, rms_norm_eps, dtype, device); } } @@ -133,7 +133,7 @@ class TextModel : public infinicore::nn::Module { protected: INFINICORE_NN_MODULE(infinicore::nn::Embedding, embed_tokens); INFINICORE_NN_MODULE_VEC(DecoderLayer, layers); - INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); + INFINICORE_NN_MODULE(Norm, norm); private: bool is_first_pp_stage() const { return pp_stage_ == 0; } diff --git a/csrc/layers/quantization/none_quantization.cpp b/csrc/layers/quantization/none_quantization.cpp index 3184da991..dde6aaedc 100644 --- a/csrc/layers/quantization/none_quantization.cpp +++ b/csrc/layers/quantization/none_quantization.cpp @@ -82,12 +82,14 @@ std::vector NoneQuantization::split_params( std::vector result; auto weight_it = params.find("weight"); auto bias_it = params.find("bias"); + const int weight_narrow_dim = + weight_prepacked_ && narrow_dim >= 0 ? 1 - narrow_dim : narrow_dim; for (const auto &s : splits) { result.push_back({s.prefix + ".weight", infinicore::nn::Parameter( - weight_it->second->narrow({{static_cast(narrow_dim), s.start, s.size}}), - narrow_dim, tp_rank, tp_size, s.num_shards)}); + weight_it->second->narrow({{static_cast(weight_narrow_dim), s.start, s.size}}), + weight_narrow_dim, tp_rank, tp_size, s.num_shards)}); if (bias_it != params.end()) { result.push_back({s.prefix + ".bias", infinicore::nn::Parameter( @@ -103,7 +105,10 @@ std::shared_ptr NoneQuantization::process_weights_after_loadin const infinicore::Device &device, int /*split_dim*/) const { - // Controlled by --pre-transpose CLI flag, default off. + // Pre-packing is opt-in. In particular, do not materialize every model + // weight on the accelerator merely because the target is Ascend: doing so + // turns model loading into a sequence of very large device-side transpose + // kernels and can stall before inference starts. if (!global_state::get_infinilm_config().pre_transpose) { return nullptr; } diff --git a/csrc/models/lfm2/lfm2_allocate_cache.cpp b/csrc/models/lfm2/lfm2_allocate_cache.cpp new file mode 100644 index 000000000..6b7ef2880 --- /dev/null +++ b/csrc/models/lfm2/lfm2_allocate_cache.cpp @@ -0,0 +1,101 @@ +#include "lfm2_allocate_cache.hpp" + +#include "../../global_state/global_state.hpp" + +#include + +#include +#include +#include +#include + +namespace infinilm::models::lfm2 { + +AllocatedLfm2Cache allocate_lfm2_cache_tensors( + const cache::CacheConfig *cache_config, + const std::shared_ptr &model_config, + backends::AttentionBackend attention_backend) { + if (cache_config == nullptr) { + return {}; + } + if (model_config == nullptr) { + throw std::runtime_error("allocate_lfm2_cache_tensors: model config is null"); + } + + const size_t num_layers = model_config->get("num_hidden_layers"); + const size_t hidden_size = model_config->get("hidden_size"); + const size_t head_dim = model_config->get("head_dim"); + const size_t num_kv_heads = model_config->get("num_key_value_heads"); + const size_t max_positions = + model_config->get("max_position_embeddings"); + const size_t state_length = model_config->get("conv_L_cache") - 1; + const auto layer_types = + model_config->get>("layer_types"); + const auto dtype = model_config->get_dtype(); + const auto kv_dtype = model_config->get_kv_cache_dtype(); + + const auto &rank_info = + infinilm::global_state::get_tensor_model_parallel_rank_info(); + const size_t pp_size = static_cast(rank_info.pp_size); + const size_t pp_stage = static_cast(rank_info.pp_stage); + const size_t local_begin = num_layers * pp_stage / pp_size; + const size_t local_end = num_layers * (pp_stage + 1) / pp_size; + + AllocatedLfm2Cache result; + result.kv_cache_tensors.resize(num_layers); + result.conv_state_tensors.resize(num_layers); + const auto device = infinicore::context::getDevice(); + + if (attention_backend == backends::AttentionBackend::STATIC_ATTN) { + auto config = dynamic_cast(cache_config); + if (config == nullptr) { + throw std::runtime_error( + "allocate_lfm2_cache_tensors: invalid static cache config"); + } + // Row 0 is immutable zero history; static scheduling uses row 1 for + // its current request. A new prefill must not read the previous + // request's terminal ShortConv state. + const size_t state_pool_size = config->max_batch_size() + 1; + for (size_t i = local_begin; i < local_end; ++i) { + if (layer_types.at(i) == "full_attention") { + result.kv_cache_tensors[i] = cache::StaticKVCache::create_layer_kv_cache( + head_dim, head_dim, num_kv_heads, num_kv_heads, + max_positions, kv_dtype, *config); + } else if (layer_types.at(i) == "short_conv") { + result.conv_state_tensors[i] = infinicore::Tensor::zeros( + {state_pool_size, hidden_size, state_length}, + dtype, device); + } + } + return result; + } + + if (attention_backend == backends::AttentionBackend::PAGED_ATTN + || attention_backend == backends::AttentionBackend::FLASH_ATTN) { + auto config = dynamic_cast(cache_config); + if (config == nullptr) { + throw std::runtime_error( + "allocate_lfm2_cache_tensors: invalid paged cache config"); + } + const size_t state_pool_size = + std::max(2, config->num_blocks() / 4); + for (size_t i = local_begin; i < local_end; ++i) { + if (layer_types.at(i) == "full_attention") { + result.kv_cache_tensors[i] = cache::PagedKVCache::create_layer_kv_cache( + head_dim, head_dim, num_kv_heads, num_kv_heads, + kv_dtype, *config); + } else if (layer_types.at(i) == "short_conv") { + result.conv_state_tensors[i] = infinicore::Tensor::zeros( + {state_pool_size, hidden_size, state_length}, + dtype, device); + } + } + infinicore::context::syncStream(); + return result; + } + + throw std::runtime_error( + "allocate_lfm2_cache_tensors: unsupported attention backend"); +} + +} // namespace infinilm::models::lfm2 diff --git a/csrc/models/lfm2/lfm2_allocate_cache.hpp b/csrc/models/lfm2/lfm2_allocate_cache.hpp new file mode 100644 index 000000000..b9132d955 --- /dev/null +++ b/csrc/models/lfm2/lfm2_allocate_cache.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include "../../backends/attention_backends.hpp" +#include "../../cache/kv_cache.hpp" +#include "../../config/model_config.hpp" + +#include + +#include +#include + +namespace infinilm::models::lfm2 { + +struct AllocatedLfm2Cache { + std::vector kv_cache_tensors; + std::vector conv_state_tensors; +}; + +AllocatedLfm2Cache allocate_lfm2_cache_tensors( + const cache::CacheConfig *cache_config, + const std::shared_ptr &model_config, + backends::AttentionBackend attention_backend); + +} // namespace infinilm::models::lfm2 diff --git a/csrc/models/lfm2/lfm2_decoder_layer.cpp b/csrc/models/lfm2/lfm2_decoder_layer.cpp new file mode 100644 index 000000000..493c833ce --- /dev/null +++ b/csrc/models/lfm2/lfm2_decoder_layer.cpp @@ -0,0 +1,72 @@ +#include "lfm2_decoder_layer.hpp" + +#include + +#include +#include + +namespace infinilm::models::lfm2 { + +Lfm2DecoderLayer::Lfm2DecoderLayer( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) + : layer_idx_(layer_idx) { + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const double rms_norm_eps = model_config->get("rms_norm_eps"); + + INFINICORE_NN_MODULE_INIT( + operator_norm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT( + ffn_norm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT(feed_forward, model_config, device); + + const auto layer_types = + model_config->get>("layer_types"); + layer_type_ = layer_types.at(layer_idx_); + if (layer_type_ == "full_attention") { + INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx_, device); + } else if (layer_type_ == "short_conv") { + INFINICORE_NN_MODULE_INIT(conv, model_config, layer_idx_, device); + } else { + throw std::runtime_error( + "Lfm2DecoderLayer: unsupported layer type '" + layer_type_ + "'"); + } +} + +std::tuple Lfm2DecoderLayer::forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) { + operator_norm_->forward_inplace(hidden_states, residual); + if (layer_type_ == "full_attention") { + hidden_states = self_attn_->forward(positions, hidden_states); + } else { + hidden_states = conv_->forward(hidden_states); + } + + ffn_norm_->forward_inplace(hidden_states, residual); + hidden_states = feed_forward_->forward(hidden_states); + return std::make_tuple(hidden_states, residual); +} + +infinicore::Tensor Lfm2DecoderLayer::forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states) { + auto residual = hidden_states; + hidden_states = operator_norm_->forward(hidden_states); + if (layer_type_ == "full_attention") { + hidden_states = self_attn_->forward(positions, hidden_states); + } else { + hidden_states = conv_->forward(hidden_states); + } + hidden_states = infinicore::op::add(residual, hidden_states); + + residual = hidden_states; + hidden_states = ffn_norm_->forward(hidden_states); + hidden_states = feed_forward_->forward(hidden_states); + return infinicore::op::add(residual, hidden_states); +} + +} // namespace infinilm::models::lfm2 diff --git a/csrc/models/lfm2/lfm2_decoder_layer.hpp b/csrc/models/lfm2/lfm2_decoder_layer.hpp new file mode 100644 index 000000000..69e42982b --- /dev/null +++ b/csrc/models/lfm2/lfm2_decoder_layer.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include "../qwen3/qwen3_attention.hpp" +#include "lfm2_mlp.hpp" +#include "lfm2_short_conv.hpp" +#include "lfm2_rms_norm.hpp" + +#include +#include + +#include +#include + +namespace infinilm::models::lfm2 { + +class Lfm2DecoderLayer : public infinicore::nn::Module { +public: + Lfm2DecoderLayer(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + std::tuple forward( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual); + + infinicore::Tensor forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states); + + size_t layer_idx() const { return layer_idx_; } + +protected: + INFINICORE_NN_MODULE(Lfm2RMSNorm, operator_norm); + INFINICORE_NN_MODULE(Lfm2RMSNorm, ffn_norm); + INFINICORE_NN_MODULE(infinilm::models::qwen3::Qwen3Attention, self_attn); + INFINICORE_NN_MODULE(Lfm2ShortConv, conv); + INFINICORE_NN_MODULE(Lfm2MLP, feed_forward); + +private: + size_t layer_idx_{0}; + std::string layer_type_; +}; + +} // namespace infinilm::models::lfm2 diff --git a/csrc/models/lfm2/lfm2_for_causal_lm.cpp b/csrc/models/lfm2/lfm2_for_causal_lm.cpp new file mode 100644 index 000000000..56d916ed7 --- /dev/null +++ b/csrc/models/lfm2/lfm2_for_causal_lm.cpp @@ -0,0 +1,108 @@ +#include "lfm2_for_causal_lm.hpp" + +#include "../../global_state/global_state.hpp" +#include "../models_registry.hpp" +#include "lfm2_allocate_cache.hpp" + +#include +#include +#include +#include +#include +#include + +namespace infinilm::models::lfm2 { + +Lfm2ForCausalLM::Lfm2ForCausalLM( + std::shared_ptr model_config, + const infinicore::Device &device) + : Lfm2CausalLMBase(std::move(model_config), device) {} + +void Lfm2ForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { + auto &forward_context = infinilm::global_state::get_forward_context(); + if (cache_config == nullptr) { + cache_config_.reset(); + forward_context.kv_cache_vec.clear(); + forward_context.conv_state_vec.clear(); + forward_context.ssm_state_vec.clear(); + return; + } + + cache_config_ = cache_config->unique_copy(); + const auto backend = + infinilm::global_state::get_infinilm_config().attention_backend; + auto allocated = allocate_lfm2_cache_tensors( + cache_config, model_config_, backend); + forward_context.kv_cache_vec = std::move(allocated.kv_cache_tensors); + forward_context.conv_state_vec = std::move(allocated.conv_state_tensors); + forward_context.ssm_state_vec.clear(); +} + +std::shared_ptr create_lfm2_model_config( + std::shared_ptr model_config) { + if (model_config->get("model_type") != "lfm2") { + throw std::runtime_error( + "create_lfm2_model_config: model_type must be 'lfm2'"); + } + + auto &json = model_config->get_config_json(); + const size_t hidden_size = json.at("hidden_size").get(); + const size_t num_heads = json.at("num_attention_heads").get(); + if (hidden_size % num_heads != 0) { + throw std::runtime_error( + "create_lfm2_model_config: hidden_size must be divisible by num_attention_heads"); + } + json["head_dim"] = hidden_size / num_heads; + + if (!json.contains("rms_norm_eps")) { + json["rms_norm_eps"] = json.value( + "norm_eps", json.value("block_norm_eps", 1e-5)); + } + + if (!json.contains("intermediate_size")) { + size_t intermediate = json.at("block_ff_dim").get(); + if (json.value("block_auto_adjust_ff_dim", false)) { + intermediate = static_cast(2.0 * intermediate / 3.0); + intermediate = static_cast( + json.value("block_ffn_dim_multiplier", 1.0) + * static_cast(intermediate)); + const size_t multiple = json.value("block_multiple_of", 1UL); + intermediate = multiple * ((intermediate + multiple - 1) / multiple); + } + json["intermediate_size"] = intermediate; + } + + if (!json.contains("layer_types")) { + const size_t num_layers = json.at("num_hidden_layers").get(); + std::unordered_set attention_layers; + for (const auto &idx : json.at("full_attn_idxs")) { + attention_layers.insert(idx.get()); + } + std::vector layer_types; + layer_types.reserve(num_layers); + for (size_t i = 0; i < num_layers; ++i) { + layer_types.push_back( + attention_layers.count(i) ? "full_attention" : "short_conv"); + } + json["layer_types"] = layer_types; + } + + if (!json.contains("rope_theta") && json.contains("rope_parameters")) { + json["rope_theta"] = json.at("rope_parameters").value( + "rope_theta", 10000.0); + } + json["attention_bias"] = json.value("attention_bias", false); + json["attention_output_bias"] = + json.value("attention_output_bias", false); + json["mlp_bias"] = json.value("mlp_bias", false); + return model_config; +} + +} // namespace infinilm::models::lfm2 + +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL( + lfm2, + infinilm::models::lfm2::Lfm2ForCausalLM, + infinilm::models::lfm2::create_lfm2_model_config); +} // namespace diff --git a/csrc/models/lfm2/lfm2_for_causal_lm.hpp b/csrc/models/lfm2/lfm2_for_causal_lm.hpp new file mode 100644 index 000000000..03cdeb99f --- /dev/null +++ b/csrc/models/lfm2/lfm2_for_causal_lm.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include "../../layers/causal_lm_templates/text_causal_lm.hpp" +#include "../../layers/causal_lm_templates/text_model.hpp" +#include "lfm2_decoder_layer.hpp" + +#include + +namespace infinilm::models::lfm2 { + +using Lfm2Model = + infinilm::layers::causal_lm_templates::TextModel; +using Lfm2CausalLMBase = + infinilm::layers::causal_lm_templates::TextCausalLM; + +class Lfm2ForCausalLM : public Lfm2CausalLMBase { +public: + Lfm2ForCausalLM( + std::shared_ptr model_config, + const infinicore::Device &device); + + void reset_cache(const cache::CacheConfig *cache_config) override; +}; + +std::shared_ptr create_lfm2_model_config( + std::shared_ptr model_config); + +} // namespace infinilm::models::lfm2 diff --git a/csrc/models/lfm2/lfm2_mlp.hpp b/csrc/models/lfm2/lfm2_mlp.hpp new file mode 100644 index 000000000..496c2b90f --- /dev/null +++ b/csrc/models/lfm2/lfm2_mlp.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include "../../layers/mlp/mlp.hpp" +#include +#include + +namespace infinilm::models::lfm2 { + +class Lfm2MLP : public infinilm::layers::mlp::MLP { +public: + using infinilm::layers::mlp::MLP::MLP; + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const { + if (hidden_states->dtype() == infinicore::DataType::F32) { + return infinilm::layers::mlp::MLP::forward(hidden_states); + } + auto input = hidden_states; + auto [gate, up] = gate_up_proj_->forward_split(input); + // Reference F.silu(gate) materializes a low-precision tensor before + // multiplication by up. A fused SwiGLU has a different rounding boundary. + auto activated = infinicore::op::silu(gate); + auto intermediate = infinicore::op::mul(activated, up); + return down_proj_->forward(intermediate); + } +}; + +} // namespace infinilm::models::lfm2 diff --git a/csrc/models/lfm2/lfm2_rms_norm.hpp b/csrc/models/lfm2/lfm2_rms_norm.hpp new file mode 100644 index 000000000..e91bb77c8 --- /dev/null +++ b/csrc/models/lfm2/lfm2_rms_norm.hpp @@ -0,0 +1,91 @@ +#pragma once + +#include +#include +#include +#include + +#include +#include +#include + +namespace infinilm::models::lfm2 { + +inline infinicore::Tensor lfm2_unit_weight( + size_t hidden, const infinicore::DataType &dtype, + const infinicore::Device &device) { + if (device.getType() == infinicore::Device::Type::ASCEND) { + // Avoid queueing aclnnInplaceOne during model construction. With a + // torch_npu-owned ACL context those tiny kernels can remain pending + // until the first RankWorker synchronization. Build the constant in + // device-pinned host memory and use InfiniCore's synchronous Ascend + // H2D path instead. + auto host = infinicore::Tensor::empty( + {hidden}, + dtype, + infinicore::Device(infinicore::Device::Type::CPU, 0), + true); + if (dtype == infinicore::DataType::F32) { + std::fill_n(reinterpret_cast(host->data()), hidden, 1.0f); + } else if (dtype == infinicore::DataType::F16) { + std::fill_n( + reinterpret_cast(host->data()), + hidden, + static_cast(0x3c00)); + } else if (dtype == infinicore::DataType::BF16) { + std::fill_n( + reinterpret_cast(host->data()), + hidden, + static_cast(0x3f80)); + } else { + throw std::runtime_error( + "LFM2 RMSNorm unit weight requires F32, F16, or BF16"); + } + return host->to(device); + } + return infinicore::Tensor::ones({hidden}, dtype, device); +} + +// Match LFM2's reference boundary: normalize in F32, cast to activation +// dtype, then multiply by the learned weight in that dtype. The unit weight +// is runtime data, not a checkpoint parameter; no ATen cast is required. +inline infinicore::Tensor lfm2_rms_norm( + const infinicore::Tensor &input, const infinicore::Tensor &weight, + float eps, const infinicore::Tensor &unit_weight) { + if (input->dtype() == infinicore::DataType::F32 + || input->device().getType() == infinicore::Device::Type::ASCEND) { + // Ascend aclnnMul does not reliably complete when the learned weight + // is exposed as a zero-stride broadcast view. Its RMSNorm backend + // already accepts the real weight and returns the requested activation + // dtype, so use that equivalent fused path on this platform. + return infinicore::op::rms_norm(input, weight, eps); + } + auto normalized = infinicore::op::rms_norm(input, unit_weight, eps); + // The final axis is the norm's feature axis. + auto strides = input->strides(); + for (auto &stride : strides) { stride = 0; } + strides.back() = 1; + auto expanded = weight->as_strided(input->shape(), strides); + return infinicore::op::mul(normalized, expanded); +} + +class Lfm2RMSNorm : public infinicore::nn::RMSNorm { +public: + Lfm2RMSNorm(size_t hidden, double eps, const infinicore::DataType &dtype, + const infinicore::Device &device) + : infinicore::nn::RMSNorm(hidden, eps, dtype, device), + unit_weight_(lfm2_unit_weight(hidden, dtype, device)) {} + + infinicore::Tensor forward(const infinicore::Tensor &input) const { + return lfm2_rms_norm(input, weight(), static_cast(eps()), unit_weight_); + } + + void forward_inplace(infinicore::Tensor &input, infinicore::Tensor &residual) const { + residual = residual ? infinicore::op::add(input, residual) : input; + input = forward(residual); + } + +private: + infinicore::Tensor unit_weight_; +}; +} // namespace infinilm::models::lfm2 diff --git a/csrc/models/lfm2/lfm2_short_conv.cpp b/csrc/models/lfm2/lfm2_short_conv.cpp new file mode 100644 index 000000000..1facac5f2 --- /dev/null +++ b/csrc/models/lfm2/lfm2_short_conv.cpp @@ -0,0 +1,221 @@ +#include "lfm2_short_conv.hpp" + +#include "../../global_state/global_state.hpp" + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace infinilm::models::lfm2 { + +Lfm2ShortConv::Lfm2ShortConv( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) + : layer_idx_(layer_idx) { + device_ = device; + hidden_size_ = model_config->get("hidden_size"); + kernel_size_ = model_config->get("conv_L_cache"); + use_bias_ = model_config->get_or("conv_bias", false); + const auto &dtype = model_config->get_dtype(); + + if (kernel_size_ < 2) { + throw std::runtime_error("Lfm2ShortConv: conv_L_cache must be at least 2"); + } + + INFINICORE_NN_MODULE_INIT( + in_proj, hidden_size_, 3 * hidden_size_, use_bias_, dtype, device); + INFINICORE_NN_MODULE_INIT( + out_proj, hidden_size_, hidden_size_, use_bias_, dtype, device); + + // Keep the nested name `conv.weight` so that the complete parameter name is + // `model.layers..conv.conv.weight`, matching the released checkpoint. + conv_weight_ = infinicore::nn::Parameter( + {hidden_size_, 1, kernel_size_}, dtype, device); + this->register_parameter("conv.weight", conv_weight_); + if (use_bias_) { + conv_bias_ = infinicore::nn::Parameter({hidden_size_}, dtype, device); + this->register_parameter("conv.bias", conv_bias_); + } +} + +infinicore::Tensor Lfm2ShortConv::causal_depthwise_conv_( + const infinicore::Tensor &input) const { + const auto &shape = input->shape(); + if (shape.size() != 3 || shape[2] != hidden_size_) { + throw std::runtime_error( + "Lfm2ShortConv: expected input shape [batch, sequence, hidden_size]"); + } + + size_t batch_size = shape[0]; + size_t sequence_length = shape[1]; + const size_t state_length = kernel_size_ - 1; + + auto &forward_context = infinilm::global_state::get_forward_context(); + auto &metadata = forward_context.mamba_metadata; + const bool has_state_routing = metadata.init_state_indices.has_value() + && metadata.final_state_indices.has_value(); + + // Paged decode flattens B one-token requests to [1, B, H]. Restore the + // logical batch dimension so every request consumes its own cache row. + infinicore::Tensor convolution_input = input; + bool restore_flattened_shape = false; + if (has_state_routing) { + const size_t request_count = metadata.init_state_indices.value()->numel(); + if (batch_size == 1 && request_count > 1) { + if (sequence_length != request_count) { + throw std::runtime_error( + "Lfm2ShortConv: packed multi-request prefill is not implemented yet"); + } + convolution_input = input->view({request_count, 1, hidden_size_}); + batch_size = request_count; + sequence_length = 1; + restore_flattened_shape = true; + } else if (request_count != batch_size) { + throw std::runtime_error( + "Lfm2ShortConv: cache index count does not match the logical batch size"); + } + } + + infinicore::Tensor state; + bool persist_state = false; + if (layer_idx_ < forward_context.conv_state_vec.size() + && forward_context.conv_state_vec[layer_idx_]) { + state = forward_context.conv_state_vec[layer_idx_]; + persist_state = true; + if (state->ndim() != 3 + || state->size(0) < batch_size + || state->size(1) != hidden_size_ + || state->size(2) != state_length) { + throw std::runtime_error("Lfm2ShortConv: incompatible convolution cache shape"); + } + } else { + state = infinicore::Tensor::zeros( + {batch_size, hidden_size_, state_length}, + input->dtype(), + input->device()); + } + + // Convert the cache from [B, H, K-1] to time-major [B, K-1, H], then + // append the current B*x values. This is the explicit left padding used by + // the reference Conv1d(padding=K-1), but it also works during token decode. + infinicore::Tensor state_batch; + if (has_state_routing) { + auto flat_state = state->view( + {state->size(0), hidden_size_ * state_length}); + state_batch = infinicore::op::embedding( + metadata.init_state_indices.value(), flat_state) + ->view({batch_size, hidden_size_, state_length}); + } else { + state_batch = state->narrow({{0, 0, batch_size}}); + } + auto history = state_batch->permute({0, 2, 1})->contiguous(); + auto combined = infinicore::Tensor::empty( + {batch_size, state_length + sequence_length, hidden_size_}, + input->dtype(), + input->device()); + combined->narrow({{1, 0, state_length}})->copy_from(history); + combined->narrow({{1, state_length, sequence_length}}) + ->copy_from(convolution_input); + infinicore::Tensor output; + const bool low_precision = input->dtype() != infinicore::DataType::F32; + if (low_precision && sequence_length > 1) { + // Conv1d prefill accumulates products in F32 and casts once. Express + // depthwise K-tap dot products as a batch of [S,K] x [K,1] GEMMs. + // This uses portable InfiniCore ops, not ATen or a CUDA-only cast. + auto windows = combined->as_strided( + {batch_size, hidden_size_, sequence_length, kernel_size_}, + {static_cast((state_length + sequence_length) * hidden_size_), + 1, static_cast(hidden_size_), static_cast(hidden_size_)}) + ->contiguous() + ->view({batch_size * hidden_size_, sequence_length, kernel_size_}); + auto kernels = conv_weight_->view({hidden_size_, kernel_size_, 1}) + ->as_strided({batch_size, hidden_size_, kernel_size_, 1}, + {0, static_cast(kernel_size_), 1, 1}) + ->contiguous() + ->view({batch_size * hidden_size_, kernel_size_, 1}); + output = infinicore::op::matmul(windows, kernels) + ->view({batch_size, hidden_size_, sequence_length}) + ->permute({0, 2, 1})->contiguous(); + } else { + auto terms = low_precision ? infinicore::Tensor::empty( + {kernel_size_, batch_size, sequence_length, hidden_size_}, + input->dtype(), input->device()) : infinicore::Tensor{}; + for (size_t kernel_idx = 0; kernel_idx < kernel_size_; ++kernel_idx) { + auto input_window = combined->narrow( + {{1, kernel_idx, sequence_length}}); + auto kernel_vector = conv_weight_ + ->narrow({{2, kernel_idx, 1}}) + ->permute({1, 2, 0}) + ->contiguous() + ->view({hidden_size_}); + // Mul does not infer a broadcasted output shape. Zero strides + // expose channel weights as [B,S,H] without B*S weight copies. + auto expanded_kernel = kernel_vector->as_strided( + input_window->shape(), {0, 0, 1}); + auto term = infinicore::op::mul(input_window, expanded_kernel); + if (low_precision) { + // Reference slow decode rounds products, then sums in F32. + terms->narrow({{0, kernel_idx, 1}}) + ->view({batch_size, sequence_length, hidden_size_})->copy_from(term); + } else { + output = output ? infinicore::op::add(output, term) : std::move(term); + } + } + if (low_precision) { + output = infinicore::op::sum(terms, {0}); + } + } + + if (use_bias_) { + auto bias = conv_bias_->view({1, 1, hidden_size_}); + output = infinicore::op::add(output, bias); + } + if (persist_state) { + auto newest_history = combined + ->narrow({{1, sequence_length, state_length}}) + ->permute({0, 2, 1}) + ->contiguous(); + if (has_state_routing) { + auto flat_state = state->view( + {state->size(0), hidden_size_ * state_length}); + auto flat_newest = newest_history->view( + {batch_size, hidden_size_ * state_length}); + infinicore::op::index_copy_( + flat_state, + flat_state, + 0, + metadata.final_state_indices.value(), + flat_newest); + } else { + state_batch->copy_from(newest_history); + } + } + if (restore_flattened_shape) { + return output->view({1, batch_size, hidden_size_}); + } + return output; +} + +infinicore::Tensor Lfm2ShortConv::forward( + const infinicore::Tensor &hidden_states) const { + auto input = hidden_states; + auto projected = in_proj_->forward(input); + + auto gate_b = projected->narrow({{2, 0, hidden_size_}}); + auto gate_c = projected->narrow({{2, hidden_size_, hidden_size_}}); + auto value = projected->narrow({{2, 2 * hidden_size_, hidden_size_}}); + + auto gated_input = infinicore::op::mul(gate_b, value); + auto convolved = causal_depthwise_conv_(gated_input); + auto gated_output = infinicore::op::mul(gate_c, convolved); + return out_proj_->forward(gated_output); +} + +} // namespace infinilm::models::lfm2 diff --git a/csrc/models/lfm2/lfm2_short_conv.hpp b/csrc/models/lfm2/lfm2_short_conv.hpp new file mode 100644 index 000000000..f12dcdb2d --- /dev/null +++ b/csrc/models/lfm2/lfm2_short_conv.hpp @@ -0,0 +1,49 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "../../layers/linear/linear.hpp" + +#include +#include +#include + +#include + +namespace infinilm::models::lfm2 { + +/** + * LFM2's gated depthwise causal convolution block. + * + * The Hugging Face reference computes + * (B, C, x) = split(in_proj(hidden_states)) + * y = out_proj(C * depthwise_causal_conv1d(B * x)) + * and keeps the last K - 1 inputs as the recurrent convolution state. + */ +class Lfm2ShortConv : public infinicore::nn::Module { +public: + Lfm2ShortConv(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + + void reset_runtime_state() const override { + in_proj_->reset_runtime_state(); + out_proj_->reset_runtime_state(); + } + +private: + infinicore::Tensor causal_depthwise_conv_(const infinicore::Tensor &input) const; + + size_t layer_idx_{0}; + size_t hidden_size_{0}; + size_t kernel_size_{0}; + bool use_bias_{false}; + + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, in_proj); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, out_proj); + INFINICORE_NN_PARAMETER(conv_weight); + INFINICORE_NN_PARAMETER(conv_bias); +}; + +} // namespace infinilm::models::lfm2 diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 117d82f9a..010dd70a7 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -57,6 +57,9 @@ def model_uses_mamba_cache(config: dict) -> bool: return ( config.get("model_type") == "mamba" or llm_config.get("model_type") == "mamba" + or config.get("model_type") == "lfm2" + or llm_config.get("model_type") == "lfm2" + or "short_conv" in layer_types or "linear_attention" in layer_types or all( key in llm_config diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index b210b0c56..971348f6d 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -741,6 +741,31 @@ def _remap_mamba(state_dict, config=None): return remapped +def _remap_lfm2(state_dict, config=None): + """Map released LFM2 parameter names to the InfiniLM module layout. + + LFM2's convolution projections already use the same names as the C++ + implementation. Only the final norm, attention norm/output names and the + three feed-forward projections need aliases. + """ + remapped = {} + replacements = ( + ("model.embedding_norm.", "model.norm."), + (".self_attn.q_layernorm.", ".self_attn.q_norm."), + (".self_attn.k_layernorm.", ".self_attn.k_norm."), + (".self_attn.out_proj.", ".self_attn.o_proj."), + (".feed_forward.w1.", ".feed_forward.gate_proj."), + (".feed_forward.w3.", ".feed_forward.up_proj."), + (".feed_forward.w2.", ".feed_forward.down_proj."), + ) + for key, tensor in state_dict.items(): + new_key = key + for source, target in replacements: + new_key = new_key.replace(source, target) + remapped[new_key] = tensor + return remapped + + def _remap_videonsa(state_dict, config=None): """Adapt VideoNSA/Qwen2.5-VL weights to the InfiniLM C++ module layout.""" key = "visual.patch_embed.proj.weight" @@ -1077,6 +1102,7 @@ def _remap_kimi_k3(state_dict, config): "baichuan": _remap_baichuan, "gpt2": _remap_gpt2, "mamba": _remap_mamba, + "lfm2": _remap_lfm2, "videonsa": _remap_videonsa, "qwen3_5": _remap_qwen3_5, "ernie4_5_moe_vl": _remap_ernie4_5_moe_vl, diff --git a/python/infinilm/processors/lfm2_processor.py b/python/infinilm/processors/lfm2_processor.py new file mode 100644 index 000000000..4248d9b08 --- /dev/null +++ b/python/infinilm/processors/lfm2_processor.py @@ -0,0 +1,76 @@ +"""Input/cache routing for LFM2's recurrent ShortConv layers.""" + +import infinicore +from typing_extensions import override + +from ..llm.scheduler import SchedulerOutput +from ..llm.static_scheduler import StaticSchedulerOutput +from .basic_llm_processor import BasicLLMProcessor +from .processor import register_processor + + +def static_short_conv_state_indices(is_prefill, prefix_hit_len, num_requests): + """Reserve row 0 for zero history and row 1 for one active request.""" + if num_requests != 1: + raise ValueError("LFM2 static scheduling currently supports one request") + if is_prefill and prefix_hit_len: + raise ValueError( + "LFM2 prefix reuse requires a matching ShortConv state snapshot" + ) + return ([0] if is_prefill else [1]), [1] + + +@register_processor("lfm2") +class Lfm2Processor(BasicLLMProcessor): + @override + def build_model_inputs( + self, + scheduler_output: SchedulerOutput | StaticSchedulerOutput, + temperature: float = 1.0, + top_p: float = 0.8, + top_k: int = 1, + **kwargs, + ) -> dict: + model_inputs = super().build_model_inputs( + scheduler_output, + temperature, + top_p, + top_k, + **kwargs, + ) + + # Static KV positions alone cannot reset recurrent convolution state. + # Route new requests from zero history and decode from the active row. + if isinstance(scheduler_output, StaticSchedulerOutput): + init_indices, final_indices = static_short_conv_state_indices( + scheduler_output.is_prefill, + scheduler_output.prefix_hit_len, + len(scheduler_output.scheduled_requests), + ) + model_inputs["mamba_init_state_indices"] = infinicore.from_list( + init_indices, dtype=infinicore.int32 + ) + model_inputs["mamba_final_state_indices"] = infinicore.from_list( + final_indices, dtype=infinicore.int32 + ) + return model_inputs + + init_indices = [] + final_indices = [] + for request in scheduler_output.scheduled_requests: + if request.mamba_cache_index is None: + raise RuntimeError( + f"Request {request.request_id} has no ShortConv cache index" + ) + init_indices.append( + 0 if scheduler_output.is_prefill else request.mamba_cache_index + ) + final_indices.append(request.mamba_cache_index) + + model_inputs["mamba_init_state_indices"] = infinicore.from_list( + init_indices, dtype=infinicore.int32 + ) + model_inputs["mamba_final_state_indices"] = infinicore.from_list( + final_indices, dtype=infinicore.int32 + ) + return model_inputs diff --git a/test/models/lfm2/native_lfm2_tiny_smoke.py b/test/models/lfm2/native_lfm2_tiny_smoke.py new file mode 100644 index 000000000..9b03b4583 --- /dev/null +++ b/test/models/lfm2/native_lfm2_tiny_smoke.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +"""Exercise native tiny LFM2 full and cached forward paths. + +This test loads the already-built InfiniCore and InfiniLM extension modules +directly. It intentionally avoids the high-level Python package so it can +isolate C++ model/runtime behavior from optional Transformers dependencies. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +from pathlib import Path +from typing import Any + +import numpy as np + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--infinicore-extension", type=Path, required=True) + parser.add_argument("--infinilm-extension", type=Path, required=True) + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--device", choices=("cpu", "cuda", "npu"), default="cpu") + parser.add_argument("--attn-backend", default="default") + parser.add_argument("--output", type=Path) + parser.add_argument("--atol", type=float, default=1e-5) + return parser.parse_args() + + +def load_extension(name: str, path: Path) -> Any: + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot create import specification for {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def tensor_to_numpy(core: Any, tensor: Any, np_dtype: Any) -> np.ndarray: + # A GPU data_ptr cannot be dereferenced by ctypes. Copy to host and finish + # the device-to-host transfer before accessing the NumPy view. + tensor = tensor.to(core.Device(core.Device.Type.CPU, 0)) + core.sync_device() + tensor = tensor.contiguous() + shape = tuple(tensor.shape) + size = int(tensor.numel()) + c_type = np.ctypeslib.as_ctypes_type(np.dtype(np_dtype)) + buffer = (c_type * size).from_address(tensor.data_ptr()) + return np.ctypeslib.as_array(buffer).copy().reshape(shape) + + +def make_input( + core: Any, + extension: Any, + token_ids: list[int], + past_length: int, + total_length: int, +) -> Any: + sequence_length = len(token_ids) + return extension.InferEngine.Input( + core.from_list([token_ids], core.DataType.I64), + position_ids=core.from_list( + [list(range(past_length, past_length + sequence_length))], + core.DataType.I64, + ), + past_sequence_lengths=core.from_list([past_length], core.DataType.I32), + total_sequence_lengths=core.from_list([total_length], core.DataType.I32), + input_offsets=core.from_list([0, sequence_length], core.DataType.I32), + cu_seqlens=core.from_list([0, total_length], core.DataType.I32), + sample_all_positions=True, + temperature=1.0, + top_k=1, + top_p=1.0, + ) + + +def main() -> None: + args = parse_args() + if args.device == "npu": + # The CANN-generated launch stubs linked into InfiniCore expect the + # Ascend runtime to be initialized before the extension is dlopen'ed. + # Importing torch_npu performs that process-wide initialization. + import torch + import torch_npu # noqa: F401 + + if not torch.npu.is_available(): + raise RuntimeError("Ascend NPU is not available") + core = load_extension("_infinicore", args.infinicore_extension) + extension = load_extension("_infinilm", args.infinilm_extension) + device_type = { + "cpu": core.Device.Type.CPU, + "cuda": core.Device.Type.NVIDIA, + "npu": core.Device.Type.ASCEND, + }[args.device] + with args.config.open("r", encoding="utf-8") as config_file: + config_text = json.dumps(json.load(config_file)) + + cache_config = extension.StaticKVCacheConfig(1, 32) + engine = extension.InferEngine( + config_text, + extension.DistConfig(1), + device_type, + cache_config, + False, + args.attn_backend, + None, + False, + "sync", + False, + ) + + rng = np.random.default_rng(20260914) + # from_blob does not own the NumPy storage. Keep every source array alive + # until load_params has copied/consumed the corresponding tensor. + source_arrays: list[np.ndarray] = [] + parameters = {} + for name, parameter in engine.state_dict()[0].items(): + shape = tuple(parameter.shape) + if name.endswith("norm.weight"): + array = np.ones(shape, dtype=np.float32) + else: + array = rng.normal(0.0, 0.02, size=shape).astype(np.float32) + source_arrays.append(array) + parameters[name] = core.from_blob( + array.ctypes.data, + list(shape), + core.DataType.F32, + core.Device(core.Device.Type.CPU, 0), + ) + + engine.load_params(parameters, True) + engine.process_weights_after_loading() + + tokens = [1, 17, 23, 5, 91, 7] + engine.reset_cache(cache_config) + full = engine.forward(make_input(core, extension, tokens, 0, len(tokens))) + full_logits = tensor_to_numpy(core, full.logits, np.float32) + + engine.reset_cache(cache_config) + engine.forward(make_input(core, extension, tokens[:-1], 0, len(tokens) - 1)) + cached = engine.forward( + make_input(core, extension, tokens[-1:], len(tokens) - 1, len(tokens)) + ) + cached_logits = tensor_to_numpy(core, cached.logits, np.float32) + + full_last = full_logits.reshape(-1, full_logits.shape[-1])[-1] + cached_last = cached_logits.reshape(-1, cached_logits.shape[-1])[-1] + max_abs_error = float(np.max(np.abs(full_last - cached_last))) + result = { + "device": args.device, + "attention_backend": args.attn_backend, + "nvidia_tf32_override": os.getenv("NVIDIA_TF32_OVERRIDE"), + "parameter_count": len(parameters), + "tokens": tokens, + "full_logits_shape": list(full_logits.shape), + "cached_logits_shape": list(cached_logits.shape), + "full_last_argmax": int(np.argmax(full_last)), + "cached_last_argmax": int(np.argmax(cached_last)), + "full_vs_cached_max_abs_error": max_abs_error, + "argmax_match": int(np.argmax(full_last)) == int(np.argmax(cached_last)), + "atol": args.atol, + } + rendered = json.dumps(result, ensure_ascii=False, indent=2) + print(rendered) + + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + + if not np.isfinite(max_abs_error): + raise AssertionError("native tiny LFM2 produced non-finite logits") + if max_abs_error > args.atol: + raise AssertionError( + f"native full and cached logits diverged: max abs error {max_abs_error}" + ) + + +if __name__ == "__main__": + main() diff --git a/test/models/lfm2/reference_lfm2.py b/test/models/lfm2/reference_lfm2.py new file mode 100644 index 000000000..b6f35e9d9 --- /dev/null +++ b/test/models/lfm2/reference_lfm2.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +"""Generate deterministic LFM2 reference values without downloading weights. + +This script intentionally uses a tiny randomly initialized model. Its purpose is +to verify model topology and cache semantics before the InfiniLM implementation is +available; it is not an accuracy test for LiquidAI/LFM2-1.2B. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F +from transformers import Lfm2Config, Lfm2ForCausalLM +from transformers.cache_utils import DynamicCache + +SEED = 20260907 +INPUT_IDS = torch.tensor([[1, 17, 23, 5, 91, 7]], dtype=torch.long) + + +def tensor_summary(tensor: torch.Tensor) -> dict[str, Any]: + value = tensor.detach().float().cpu() + flat = value.flatten() + return { + "shape": list(value.shape), + "mean": float(value.mean()), + "std": float(value.std(unbiased=False)), + "min": float(value.min()), + "max": float(value.max()), + "first_values": [float(x) for x in flat[:8]], + } + + +def first_tensor(output: Any) -> torch.Tensor: + if isinstance(output, torch.Tensor): + return output + if isinstance(output, (tuple, list)): + for item in output: + if isinstance(item, torch.Tensor): + return item + raise TypeError(f"No tensor found in hook output of type {type(output)!r}") + + +def describe_cache(cache: DynamicCache) -> list[dict[str, Any]]: + description = [] + for layer_idx, layer in enumerate(cache.layers): + item: dict[str, Any] = { + "layer_idx": layer_idx, + "cache_class": type(layer).__name__, + } + if hasattr(layer, "conv_states"): + state = layer.conv_states[0] + item["conv_state_shape"] = list(state.shape) + item["has_previous_state"] = bool(layer.has_previous_state[0]) + else: + item["key_shape"] = list(layer.keys.shape) + item["value_shape"] = list(layer.values.shape) + description.append(item) + return description + + +def verify_short_conv( + model: Lfm2ForCausalLM, + input_ids: torch.Tensor, +) -> float: + """Rebuild layer-0 ShortConv from primitive PyTorch operations.""" + layer = model.model.layers[0] + hidden = model.model.embed_tokens(input_ids) + normalized = layer.operator_norm(hidden) + + projected = layer.conv.in_proj(normalized).transpose(-1, -2) + b_gate, c_gate, x_value = projected.chunk(3, dim=-2) + bx = b_gate * x_value + + conv_out = F.conv1d( + bx, + layer.conv.conv.weight, + layer.conv.conv.bias, + padding=layer.conv.L_cache - 1, + groups=model.config.hidden_size, + )[..., : input_ids.shape[1]] + manual = layer.conv.out_proj((c_gate * conv_out).transpose(-1, -2).contiguous()) + reference = layer.conv(normalized, attention_mask=torch.ones_like(input_ids)) + return float((manual - reference).detach().abs().max()) + + +def run(config_path: Path) -> dict[str, Any]: + torch.manual_seed(SEED) + torch.use_deterministic_algorithms(True) + + config = Lfm2Config.from_json_file(str(config_path)) + model = Lfm2ForCausalLM(config).float().eval() + first_attention_idx = config.layer_types.index("full_attention") + + captured: dict[str, dict[str, Any]] = {} + + def capture(name: str): + def hook(_module, _inputs, output): + captured[name] = tensor_summary(first_tensor(output)) + + return hook + + handles = [ + model.model.layers[0].conv.register_forward_hook(capture("layer0_short_conv")), + model.model.layers[first_attention_idx].self_attn.register_forward_hook( + capture(f"layer{first_attention_idx}_attention") + ), + model.model.layers[0].feed_forward.register_forward_hook(capture("layer0_mlp")), + model.model.embedding_norm.register_forward_hook(capture("final_rms_norm")), + ] + + attention_mask = torch.ones_like(INPUT_IDS) + with torch.inference_mode(): + full_logits = model( + input_ids=INPUT_IDS, + attention_mask=attention_mask, + use_cache=False, + ).logits + + for handle in handles: + handle.remove() + + with torch.inference_mode(): + cache = DynamicCache(config=config) + prefill_logits = model( + input_ids=INPUT_IDS[:, :-1], + attention_mask=attention_mask[:, :-1], + past_key_values=cache, + use_cache=True, + ).logits + decode_logits = model( + input_ids=INPUT_IDS[:, -1:], + attention_mask=attention_mask, + past_key_values=cache, + use_cache=True, + ).logits + + cache_description = describe_cache(cache) + cache.reset() + reset_flags = [ + bool(layer.has_previous_state[0]) + for layer in cache.layers + if hasattr(layer, "has_previous_state") + ] + + cached_error = float((full_logits[:, -1] - decode_logits[:, -1]).abs().max()) + short_conv_error = verify_short_conv(model, INPUT_IDS) + + torch.testing.assert_close( + decode_logits[:, -1], + full_logits[:, -1], + rtol=1e-5, + atol=1e-6, + ) + if short_conv_error > 1e-7: + raise AssertionError(f"ShortConv reconstruction error is {short_conv_error}") + if any(reset_flags): + raise AssertionError("cache.reset() left a convolution state active") + + return { + "purpose": "LFM2 topology and cache reference; not pretrained-model accuracy", + "seed": SEED, + "input_ids": INPUT_IDS.tolist(), + "config": { + "hidden_size": config.hidden_size, + "intermediate_size_before_auto_adjust": config.intermediate_size, + "actual_mlp_intermediate_size": model.model.layers[ + 0 + ].feed_forward.w1.out_features, + "num_hidden_layers": config.num_hidden_layers, + "layer_types": config.layer_types, + "num_attention_heads": config.num_attention_heads, + "num_key_value_heads": config.num_key_value_heads, + "head_dim": model.model.layers[first_attention_idx].self_attn.head_dim, + "conv_L_cache": config.conv_L_cache, + }, + "parameter_count": sum(parameter.numel() for parameter in model.parameters()), + "state_dict_shapes": { + name: list(tensor.shape) for name, tensor in model.state_dict().items() + }, + "forward": { + "full_logits": tensor_summary(full_logits), + "prefill_logits_shape": list(prefill_logits.shape), + "decode_logits_shape": list(decode_logits.shape), + "full_last_token_argmax": int(full_logits[:, -1].argmax(dim=-1).item()), + "cached_last_token_argmax": int(decode_logits[:, -1].argmax(dim=-1).item()), + "full_vs_cached_max_abs_error": cached_error, + "short_conv_manual_max_abs_error": short_conv_error, + }, + "captured_modules": captured, + "cache_after_prefill_and_decode": cache_description, + "conv_cache_flags_after_reset": reset_flags, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--config", + type=Path, + default=Path(__file__).with_name("tiny_config.json"), + ) + parser.add_argument( + "--output", + type=Path, + help="Optional JSON output path. The result is always printed as well.", + ) + args = parser.parse_args() + + result = run(args.config) + rendered = json.dumps(result, ensure_ascii=False, indent=2) + print(rendered) + if args.output is not None: + args.output.write_text(rendered + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/test/models/lfm2/reference_lfm2_real.py b/test/models/lfm2/reference_lfm2_real.py new file mode 100644 index 000000000..bb6fce106 --- /dev/null +++ b/test/models/lfm2/reference_lfm2_real.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +"""Record a reproducible Transformers reference for LFM2-1.2B. + +Run this on the NVIDIA instance after the official weights are available. The +script performs greedy generation and writes token IDs plus per-step top-k +scores, which are more useful for backend comparison than decoded text alone. +""" + +from __future__ import annotations + +import argparse +import json +import platform +from pathlib import Path + +import torch +import transformers +from transformers import AutoModelForCausalLM, AutoTokenizer + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + required=True, + help="Local LFM2-1.2B directory or Hugging Face model ID.", + ) + parser.add_argument("--prompt", default="Who are you?") + parser.add_argument("--max-new-tokens", type=int, default=16) + parser.add_argument("--top-k", type=int, default=5) + parser.add_argument("--device", default="cuda") + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.device.startswith("cuda") and not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested, but torch.cuda.is_available() is false") + + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=False) + model = AutoModelForCausalLM.from_pretrained( + args.model, + dtype=torch.bfloat16, + attn_implementation="eager", + trust_remote_code=False, + ).to(args.device) + model.eval() + + messages = [{"role": "user", "content": args.prompt}] + encoded = tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_tensors="pt", + return_dict=True, + ) + encoded = {name: value.to(args.device) for name, value in encoded.items()} + + with torch.inference_mode(): + generated = model.generate( + **encoded, + max_new_tokens=args.max_new_tokens, + do_sample=False, + use_cache=True, + return_dict_in_generate=True, + output_scores=True, + ) + + prompt_length = encoded["input_ids"].shape[1] + sequence = generated.sequences[0].detach().cpu() + generated_ids = sequence[prompt_length:] + + step_scores = [] + for step, logits in enumerate(generated.scores): + values, indices = torch.topk(logits[0].float(), k=args.top_k) + step_scores.append( + { + "step": step, + "selected_token_id": int(generated_ids[step]), + "top_token_ids": [int(x) for x in indices.detach().cpu()], + "top_scores": [float(x) for x in values.detach().cpu()], + } + ) + + result = { + "model": args.model, + "prompt": args.prompt, + "chat_messages": messages, + "input_ids": encoded["input_ids"][0].detach().cpu().tolist(), + "generated_token_ids": generated_ids.tolist(), + "full_token_ids": sequence.tolist(), + "decoded_text": tokenizer.decode( + sequence, + skip_special_tokens=False, + clean_up_tokenization_spaces=False, + ), + "generated_text": tokenizer.decode( + generated_ids, + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + ), + "step_scores": step_scores, + "environment": { + "python": platform.python_version(), + "pytorch": torch.__version__, + "transformers": transformers.__version__, + "device": args.device, + "cuda_available": torch.cuda.is_available(), + "cuda_runtime": torch.version.cuda, + "gpu": ( + torch.cuda.get_device_name(torch.device(args.device)) + if args.device.startswith("cuda") + else None + ), + "dtype": "bfloat16", + "attention_implementation": "eager", + "decoding": "greedy", + }, + } + + rendered = json.dumps(result, ensure_ascii=False, indent=2) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + print(rendered) + + +if __name__ == "__main__": + main() diff --git a/test/models/lfm2/run_infinilm_lfm2_real.py b/test/models/lfm2/run_infinilm_lfm2_real.py new file mode 100644 index 000000000..4be131757 --- /dev/null +++ b/test/models/lfm2/run_infinilm_lfm2_real.py @@ -0,0 +1,295 @@ +#!/usr/bin/env python3 +"""Run reproducible LFM2-1.2B inference with InfiniLM. + +The script is intended for the NVIDIA and Ascend validation machines. It keeps +sampling greedy, records exact prompt/generated token IDs, and repeats the same +prompt set in one engine process so stale KV/ShortConv state is easy to detect. + +Example: + python test/models/lfm2/run_infinilm_lfm2_real.py \ + --model /data/models/LFM2-1.2B \ + --device cuda \ + --output artifacts/lfm2_infinilm_cuda.json + +To turn a Transformers reference artifact into an exact correctness check: + python test/models/lfm2/run_infinilm_lfm2_real.py \ + --model /data/models/LFM2-1.2B \ + --device cuda \ + --prompt "Who are you?" \ + --reference artifacts/lfm2_transformers_cuda.json \ + --output artifacts/lfm2_infinilm_cuda_checked.json +""" + +from __future__ import annotations + +import argparse +import json +import platform +import time +from pathlib import Path +from typing import Any + +DEFAULT_PROMPTS = ( + "Who are you?", + "请用一句中文介绍你自己。", + "Explain in three short points why recurrent state can reduce decoding work.", +) + + +def resolve_attention_backend(cache_type: str, requested: str) -> str: + """Keep cache layout consistent with the C++ attention implementation.""" + if cache_type not in ("static", "paged"): + raise ValueError(f"Unsupported cache type: {cache_type}") + if requested == "default": + return "static-attn" if cache_type == "static" else "paged-attn" + if requested not in ("static-attn", "paged-attn", "flash-attn"): + raise ValueError(f"Unsupported LFM2 attention backend: {requested}") + if cache_type == "static" and requested != "static-attn": + raise ValueError("Static cache requires the static-attn backend") + if cache_type == "paged" and requested == "static-attn": + raise ValueError("Paged cache cannot use the static-attn backend") + return requested + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True, help="Local LFM2-1.2B directory") + parser.add_argument( + "--device", + choices=("cpu", "cuda", "npu"), + default="cuda", + help="InfiniCore device: cuda=NVIDIA, npu=Ascend", + ) + parser.add_argument( + "--prompt", + action="append", + help="Prompt to run; repeat this option for multiple prompts", + ) + parser.add_argument("--max-new-tokens", type=int, default=16) + parser.add_argument( + "--repeat", + type=int, + default=2, + help="Run the whole prompt set this many times in one engine process", + ) + parser.add_argument("--cache-type", choices=("paged", "static"), default="paged") + parser.add_argument("--num-blocks", type=int, default=64) + parser.add_argument("--block-size", type=int, default=64) + parser.add_argument("--max-cache-len", type=int, default=4096) + parser.add_argument( + "--attn-backend", + default="default", + help="default selects static-attn/paged-attn to match --cache-type", + ) + parser.add_argument( + "--reference", + type=Path, + help="JSON created by reference_lfm2_real.py; its matching prompt is checked", + ) + parser.add_argument( + "--extra-reference", + action="append", + type=Path, + default=[], + help="Additional reference JSONs; every matching prompt is checked", + ) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + if args.max_new_tokens < 1: + parser.error("--max-new-tokens must be positive") + if args.repeat < 1: + parser.error("--repeat must be positive") + try: + args.attn_backend = resolve_attention_backend( + args.cache_type, args.attn_backend + ) + except ValueError as error: + parser.error(str(error)) + return args + + +def _enum_value(value: Any) -> Any: + return getattr(value, "value", value) + + +def _environment(device: str) -> dict[str, Any]: + import torch + + result: dict[str, Any] = { + "python": platform.python_version(), + "platform": platform.platform(), + "pytorch": torch.__version__, + "device": device, + "cuda_available": torch.cuda.is_available(), + "cuda_runtime": torch.version.cuda, + } + if device == "cuda" and torch.cuda.is_available(): + result["accelerator"] = torch.cuda.get_device_name(0) + elif device == "npu": + npu = getattr(torch, "npu", None) + if npu is not None and npu.is_available(): + result["accelerator"] = npu.get_device_name(0) + return result + + +def _load_reference(path: Path | None) -> dict[str, Any] | None: + if path is None: + return None + with path.open("r", encoding="utf-8") as file: + return json.load(file) + + +def _compare_reference( + reference: dict[str, Any] | None, runs: list[dict[str, Any]] +) -> dict[str, Any] | None: + if reference is None: + return None + + prompt = reference.get("prompt") + candidate = next((run for run in runs if run["prompt"] == prompt), None) + if candidate is None: + return { + "passed": False, + "reason": f"reference prompt was not run: {prompt!r}", + } + + expected_prompt_ids = reference.get("input_ids") + expected_generated_ids = reference.get("generated_token_ids") + prompt_match = candidate["prompt_token_ids"] == expected_prompt_ids + generated_match = candidate["generated_token_ids"] == expected_generated_ids + return { + "passed": prompt_match and generated_match, + "prompt": prompt, + "prompt_token_ids_match": prompt_match, + "generated_token_ids_match": generated_match, + "expected_generated_token_ids": expected_generated_ids, + "actual_generated_token_ids": candidate["generated_token_ids"], + } + + +def main() -> None: + args = parse_args() + + if args.device == "npu": + # Initialize CANN before importing the native InfiniLM extension. The + # AscendC launch stubs linked through InfiniCore are initialized at + # dlopen time and require torch_npu's process-wide runtime setup. + import torch + import torch_npu # noqa: F401 + + if not torch.npu.is_available(): + raise RuntimeError("Ascend NPU is not available") + + from infinilm import LLM, SamplingParams + + prompts = args.prompt or list(DEFAULT_PROMPTS) + sampling = SamplingParams( + temperature=1.0, + top_k=1, + top_p=1.0, + max_tokens=args.max_new_tokens, + ) + model = LLM( + model_path=args.model, + device=args.device, + tensor_parallel_size=1, + cache_type=args.cache_type, + max_batch_size=1, + max_tokens=args.max_new_tokens, + num_blocks=args.num_blocks, + block_size=args.block_size, + max_cache_len=args.max_cache_len, + temperature=1.0, + top_k=1, + top_p=1.0, + enable_graph=False, + attn_backend=args.attn_backend, + enable_prefix_caching=False, + ) + + runs: list[dict[str, Any]] = [] + try: + # The repetition loop is outside the prompt loop, giving A/B/C/A/B/C. + # A repeated result must therefore survive both row reuse and intervening + # requests without stale KV or ShortConv state contamination. + for repetition in range(args.repeat): + for prompt_index, prompt in enumerate(prompts): + messages = [{"role": "user", "content": prompt}] + start = time.perf_counter() + request = model.chat( + messages=messages, + sampling_params=sampling, + use_tqdm=False, + )[0] + elapsed = time.perf_counter() - start + completion = request.outputs[0] + run = { + "repetition": repetition, + "prompt_index": prompt_index, + "prompt": prompt, + "prompt_token_ids": request.prompt_token_ids, + "generated_token_ids": completion.token_ids, + "generated_text": completion.text, + "finish_reason": _enum_value(completion.finish_reason), + "elapsed_seconds": elapsed, + "generated_tokens_per_second": ( + len(completion.token_ids) / elapsed if elapsed > 0 else None + ), + } + runs.append(run) + print(json.dumps(run, ensure_ascii=False)) + finally: + model.close() + + sequences_by_prompt: dict[int, list[list[int]]] = {} + for run in runs: + sequences_by_prompt.setdefault(run["prompt_index"], []).append( + run["generated_token_ids"] + ) + repeat_consistent = all( + all(sequence == sequences[0] for sequence in sequences[1:]) + for sequences in sequences_by_prompt.values() + ) + + reference = _load_reference(args.reference) + reference_comparison = _compare_reference(reference, runs) + reference_comparisons = [reference_comparison] if reference_comparison else [] + reference_comparisons.extend( + _compare_reference(_load_reference(path), runs) for path in args.extra_reference + ) + result = { + "model": args.model, + "configuration": { + "cache_type": args.cache_type, + "num_blocks": args.num_blocks, + "block_size": args.block_size, + "max_cache_len": args.max_cache_len, + "attention_backend": args.attn_backend, + "max_new_tokens": args.max_new_tokens, + "repeat": args.repeat, + "decoding": "greedy", + "dtype": "from model config", + }, + "environment": _environment(args.device), + "repeat_consistent": repeat_consistent, + "reference_comparison": reference_comparison, + "reference_comparisons": reference_comparisons, + "runs": runs, + } + + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps(result, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + print(json.dumps(result, ensure_ascii=False, indent=2)) + + if not repeat_consistent: + raise SystemExit("Repeated greedy outputs differ; cache state may be stale") + if any(not comparison["passed"] for comparison in reference_comparisons): + raise SystemExit("InfiniLM token IDs differ from the Transformers reference") + + +if __name__ == "__main__": + main() diff --git a/test/models/lfm2/test_low_precision_contract.py b/test/models/lfm2/test_low_precision_contract.py new file mode 100644 index 000000000..11e7d98db --- /dev/null +++ b/test/models/lfm2/test_low_precision_contract.py @@ -0,0 +1,61 @@ +"""Reference rounding boundaries, independent of GPU and native extensions.""" + +import unittest + +import torch +import torch.nn.functional as F + + +class Lfm2LowPrecisionContractTest(unittest.TestCase): + def setUp(self): + torch.manual_seed(20260915) + + def test_rms_norm_casts_before_learned_weight(self): + value = torch.randn(2, 5, 16).bfloat16() + weight = torch.randn(16).bfloat16() + normalized = value.float() * torch.rsqrt( + value.float().square().mean(-1, keepdim=True) + 1e-5 + ) + expected = normalized.bfloat16() * weight + unit_norm = (normalized * torch.ones_like(weight).float()).bfloat16() + torch.testing.assert_close(unit_norm * weight, expected, rtol=0, atol=0) + + def test_mlp_materializes_silu_before_product(self): + gate, up = torch.randn(2, 5, 16).bfloat16(), torch.randn(2, 5, 16).bfloat16() + actual = F.silu(gate.float()).bfloat16() * up + expected = F.silu(gate) * up + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_neox_rotary_rounds_tables_and_both_products(self): + value = torch.randn(2, 5, 16).bfloat16() + cosine, sine = torch.randn(5, 8).bfloat16(), torch.randn(5, 8).bfloat16() + first, second = value.chunk(2, -1) + actual = torch.cat( + (first * cosine - second * sine, second * cosine + first * sine), -1 + ) + rotated_half = torch.cat((-second, first), -1) + expected = value * torch.cat((cosine, cosine), -1) + rotated_half * torch.cat( + (sine, sine), -1 + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_python_scalar_keeps_float32_opmath(self): + value = torch.arange(1, 1000).bfloat16() + alpha = 8**-0.5 + expected = (value.float() * alpha).bfloat16() + torch.testing.assert_close(value * alpha, expected, rtol=0, atol=0) + rounded_alpha = value * torch.tensor(alpha).bfloat16() + self.assertGreater(int(torch.count_nonzero(rounded_alpha != expected)), 0) + + def test_attention_rounds_qk_before_scale(self): + query, key = torch.randn(2, 5, 8).bfloat16(), torch.randn(2, 8, 5).bfloat16() + alpha = 8**-0.5 + expected = (query.float() @ key.float()).bfloat16() * alpha + actual = (query @ key) * alpha + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + fused_scale = ((query.float() @ key.float()) * alpha).bfloat16() + self.assertGreater(int(torch.count_nonzero(fused_scale != expected)), 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/lfm2/test_run_config.py b/test/models/lfm2/test_run_config.py new file mode 100644 index 000000000..503ce9df0 --- /dev/null +++ b/test/models/lfm2/test_run_config.py @@ -0,0 +1,38 @@ +import unittest + +from run_infinilm_lfm2_real import _compare_reference, resolve_attention_backend + + +class Lfm2RunConfigTest(unittest.TestCase): + def test_default_backend_matches_cache_layout(self): + self.assertEqual(resolve_attention_backend("static", "default"), "static-attn") + self.assertEqual(resolve_attention_backend("paged", "default"), "paged-attn") + self.assertEqual(resolve_attention_backend("paged", "flash-attn"), "flash-attn") + + def test_invalid_backend_pairs_are_rejected(self): + for cache_type, backend in ( + ("paged", "static-attn"), + ("static", "paged-attn"), + ("static", "flash-attn"), + ("paged", "unknown"), + ("unknown", "default"), + ): + with self.subTest(cache_type=cache_type, backend=backend): + with self.assertRaises(ValueError): + resolve_attention_backend(cache_type, backend) + + def test_each_reference_prompt_is_checked_independently(self): + runs = [ + {"prompt": "A", "prompt_token_ids": [1, 2], "generated_token_ids": [3, 4]}, + {"prompt": "B", "prompt_token_ids": [1, 5], "generated_token_ids": [6, 7]}, + ] + good = {"prompt": "A", "input_ids": [1, 2], "generated_token_ids": [3, 4]} + bad = {"prompt": "B", "input_ids": [1, 5], "generated_token_ids": [6, 8]} + self.assertTrue(_compare_reference(good, runs)["passed"]) + self.assertFalse(_compare_reference(bad, runs)["passed"]) + self.assertFalse(_compare_reference({"prompt": "missing"}, runs)["passed"]) + self.assertIsNone(_compare_reference(None, runs)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/lfm2/test_short_conv_contract.py b/test/models/lfm2/test_short_conv_contract.py new file mode 100644 index 000000000..2cab9fda2 --- /dev/null +++ b/test/models/lfm2/test_short_conv_contract.py @@ -0,0 +1,125 @@ +"""Mathematical contract for the explicit LFM2 ShortConv decomposition.""" + +from __future__ import annotations + +import unittest + +import torch +import torch.nn.functional as F + + +def stateful_depthwise_conv( + values: torch.Tensor, + weight: torch.Tensor, + state: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Mirror the C++ implementation using [B, H, K-1] recurrent state.""" + kernel_size = weight.shape[-1] + history = state.transpose(1, 2) + combined = torch.cat((history, values), dim=1) + + if values.dtype != torch.float32 and values.shape[1] > 1: + batch, sequence, hidden = values.shape + windows = ( + combined.as_strided( + (batch, hidden, sequence, kernel_size), + ((sequence + kernel_size - 1) * hidden, 1, hidden, hidden), + ) + .contiguous() + .view(batch * hidden, sequence, kernel_size) + ) + kernels = ( + weight.view(hidden, kernel_size, 1) + .unsqueeze(0) + .expand(batch, hidden, kernel_size, 1) + .contiguous() + .view(batch * hidden, kernel_size, 1) + ) + output = torch.bmm(windows.float(), kernels.float()).to(values.dtype) + output = output.view(batch, hidden, sequence).transpose(1, 2).contiguous() + return output, combined[:, -kernel_size + 1 :, :].transpose(1, 2) + + output = None + terms = [] + for kernel_idx in range(kernel_size): + window = combined[:, kernel_idx : kernel_idx + values.shape[1], :] + term = window * weight[:, 0, kernel_idx].view(1, 1, -1) + terms.append(term) + output = term if output is None else output + term + + if values.dtype != torch.float32: + output = torch.stack(terms).float().sum(dim=0).to(values.dtype) + + newest_state = combined[:, -kernel_size + 1 :, :].transpose(1, 2) + return output, newest_state + + +class Lfm2ShortConvContractTest(unittest.TestCase): + def setUp(self): + torch.manual_seed(20260914) + self.batch = 2 + self.sequence = 7 + self.hidden = 5 + self.kernel = 3 + self.values = torch.randn(self.batch, self.sequence, self.hidden) + self.weight = torch.randn(self.hidden, 1, self.kernel) + + def test_prefill_matches_grouped_conv1d(self): + state = torch.zeros(self.batch, self.hidden, self.kernel - 1) + actual, _ = stateful_depthwise_conv(self.values, self.weight, state) + expected = F.conv1d( + self.values.transpose(1, 2), + self.weight, + padding=self.kernel - 1, + groups=self.hidden, + )[..., : self.sequence].transpose(1, 2) + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-6) + + def test_prefill_then_decode_matches_full_sequence(self): + split = self.sequence - 1 + state = torch.zeros(self.batch, self.hidden, self.kernel - 1) + prefill, state = stateful_depthwise_conv( + self.values[:, :split], self.weight, state + ) + decode, state = stateful_depthwise_conv( + self.values[:, split:], self.weight, state + ) + cached = torch.cat((prefill, decode), dim=1) + + full, expected_state = stateful_depthwise_conv( + self.values, + self.weight, + torch.zeros_like(state), + ) + torch.testing.assert_close(cached, full, rtol=1e-6, atol=1e-6) + torch.testing.assert_close(state, expected_state, rtol=0, atol=0) + + def test_bf16_prefill_has_single_output_rounding(self): + values, weight = self.values.bfloat16(), self.weight.bfloat16() + state = torch.zeros( + self.batch, self.hidden, self.kernel - 1, dtype=torch.bfloat16 + ) + actual, _ = stateful_depthwise_conv(values, weight, state) + expected = F.conv1d( + values.float().transpose(1, 2), + weight.float(), + padding=self.kernel - 1, + groups=self.hidden, + )[..., : self.sequence] + torch.testing.assert_close( + actual, expected.transpose(1, 2).bfloat16(), rtol=0, atol=0 + ) + + def test_bf16_decode_sums_rounded_products_in_float32(self): + values, weight = self.values[:, -1:].bfloat16(), self.weight.bfloat16() + state = torch.randn(self.batch, self.hidden, self.kernel - 1).bfloat16() + actual, newest = stateful_depthwise_conv(values, weight, state) + combined = torch.cat((state, values.transpose(1, 2)), dim=-1) + products = combined * weight[:, 0, :] + expected = products.float().sum(-1).bfloat16().unsqueeze(1) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(newest, combined[..., 1:], rtol=0, atol=0) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/lfm2/test_state_routing.py b/test/models/lfm2/test_state_routing.py new file mode 100644 index 000000000..d8a875d29 --- /dev/null +++ b/test/models/lfm2/test_state_routing.py @@ -0,0 +1,40 @@ +"""Test the pure routing helper without requiring a compiled GPU extension.""" + +import ast +import unittest +from pathlib import Path + +source = ( + Path(__file__).resolve().parents[3] / "python/infinilm/processors/lfm2_processor.py" +) +tree = ast.parse(source.read_text(encoding="utf-8")) +function = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) + and node.name == "static_short_conv_state_indices" +) +namespace = {} +exec( + compile(ast.Module(body=[function], type_ignores=[]), str(source), "exec"), + namespace, +) +route = namespace["static_short_conv_state_indices"] + + +class Lfm2StateRoutingTest(unittest.TestCase): + def test_new_requests_never_read_previous_terminal_state(self): + # A prefill, A decode, B prefill, B decode, A prefill. + self.assertEqual( + [route(stage, 0, 1) for stage in (True, False, True, False, True)], + [([0], [1]), ([1], [1]), ([0], [1]), ([1], [1]), ([0], [1])], + ) + + def test_unimplemented_prefix_and_batch_modes_are_rejected(self): + for arguments in ((True, 5, 1), (True, 0, 0), (True, 0, 2)): + with self.subTest(arguments=arguments), self.assertRaises(ValueError): + route(*arguments) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/lfm2/test_weight_remap.py b/test/models/lfm2/test_weight_remap.py new file mode 100644 index 000000000..d8246bf46 --- /dev/null +++ b/test/models/lfm2/test_weight_remap.py @@ -0,0 +1,93 @@ +"""Unit tests for the LFM2 checkpoint-name adapter. + +This test deliberately stubs the compiled ``infinicore`` module so the name +mapping can be checked before a native InfiniLM build is available. +""" + +from __future__ import annotations + +import importlib.util +import sys +import types +import unittest +from pathlib import Path + +import torch + + +def _load_modeling_utils(): + fake_infinicore = types.ModuleType("infinicore") + fake_infinicore.float32 = object() + fake_infinicore.float16 = object() + fake_infinicore.bfloat16 = object() + fake_infinicore.int8 = object() + fake_infinicore.int32 = object() + fake_infinicore.int64 = object() + fake_infinicore.dtype = object() + fake_infinicore.device = object() + fake_infinicore.Tensor = object + fake_infinicore.nn = types.SimpleNamespace(Module=object) + sys.modules.setdefault("infinicore", fake_infinicore) + + module_path = ( + Path(__file__).resolve().parents[3] + / "python" + / "infinilm" + / "modeling_utils.py" + ) + spec = importlib.util.spec_from_file_location( + "infinilm_lfm2_modeling_utils", module_path + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +class Lfm2WeightRemapTest(unittest.TestCase): + def test_all_released_name_patterns(self): + module = _load_modeling_utils() + tensor = torch.zeros(1) + state_dict = { + "model.embedding_norm.weight": tensor, + "model.layers.2.self_attn.q_layernorm.weight": tensor, + "model.layers.2.self_attn.k_layernorm.weight": tensor, + "model.layers.2.self_attn.out_proj.weight": tensor, + "model.layers.0.feed_forward.w1.weight": tensor, + "model.layers.0.feed_forward.w2.weight": tensor, + "model.layers.0.feed_forward.w3.weight": tensor, + "model.layers.0.conv.conv.weight": tensor, + "model.layers.0.conv.in_proj.weight": tensor, + "model.layers.0.conv.out_proj.weight": tensor, + "model.layers.0.operator_norm.weight": tensor, + "model.layers.0.ffn_norm.weight": tensor, + "model.embed_tokens.weight": tensor, + } + + remapped = module._remap_lfm2(state_dict) + self.assertEqual( + set(remapped), + { + "model.norm.weight", + "model.layers.2.self_attn.q_norm.weight", + "model.layers.2.self_attn.k_norm.weight", + "model.layers.2.self_attn.o_proj.weight", + "model.layers.0.feed_forward.gate_proj.weight", + "model.layers.0.feed_forward.down_proj.weight", + "model.layers.0.feed_forward.up_proj.weight", + "model.layers.0.conv.conv.weight", + "model.layers.0.conv.in_proj.weight", + "model.layers.0.conv.out_proj.weight", + "model.layers.0.operator_norm.weight", + "model.layers.0.ffn_norm.weight", + "model.embed_tokens.weight", + }, + ) + + def test_registry_selects_lfm2_remapper(self): + module = _load_modeling_utils() + self.assertIs(module._WEIGHT_REMAPPER["lfm2"], module._remap_lfm2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/lfm2/tiny_config.json b/test/models/lfm2/tiny_config.json new file mode 100644 index 000000000..93b05834b --- /dev/null +++ b/test/models/lfm2/tiny_config.json @@ -0,0 +1,35 @@ +{ + "architectures": [ + "Lfm2ForCausalLM" + ], + "model_type": "lfm2", + "dtype": "float32", + "vocab_size": 128, + "hidden_size": 32, + "block_dim": 32, + "block_ff_dim": 96, + "block_ffn_dim_multiplier": 1.0, + "block_multiple_of": 16, + "block_auto_adjust_ff_dim": true, + "num_hidden_layers": 8, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "full_attn_idxs": [ + 2, + 5, + 7 + ], + "max_position_embeddings": 128, + "norm_eps": 1e-05, + "conv_L_cache": 3, + "conv_bias": false, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 1000000.0 + }, + "bos_token_id": 1, + "eos_token_id": 2, + "pad_token_id": 0, + "tie_word_embeddings": true, + "use_cache": true +} diff --git a/test/models/lfm2/validate_lfm2_cache_and_timing.py b/test/models/lfm2/validate_lfm2_cache_and_timing.py new file mode 100644 index 000000000..e84796a51 --- /dev/null +++ b/test/models/lfm2/validate_lfm2_cache_and_timing.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +"""Validate recurrent cache reuse and record synchronized model-forward timing. + +F32/static is a numerical correctness gate. BF16 full/cache comparisons are +diagnostics because GEMM and attention shapes change rounding. All modes gate +finite logits and repeatable generation. Timings exclude tokenization/loading, +include engine dispatch and device synchronization, and are not optimization data. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import tempfile +import time +from pathlib import Path + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--reference-dir", type=Path, required=True) + parser.add_argument("--dtype", choices=("float32", "bfloat16"), default="bfloat16") + parser.add_argument("--cache-type", choices=("static", "paged"), default="paged") + parser.add_argument("--steps", type=int, default=64) + parser.add_argument("--compare-full-steps", type=int, default=8) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.steps < 1 or args.compare_full_steps < 0: + parser.error( + "steps must be positive; compare-full-steps must be nonnegative (0 = timing only)" + ) + if args.dtype == "float32" and args.cache_type == "paged": + parser.error("NVIDIA Paged F32 prefill is unsupported") + + import infinicore + import numpy as np + import torch + from infinicore.lib import _infinicore as core + from infinilm.infer_engine import InferEngine + from infinilm.lib import _infinilm + from infinilm.modeling_utils import load_model_state_dict_by_file + from native_lfm2_tiny_smoke import tensor_to_numpy + + references = [ + json.loads((args.reference_dir / name).read_text(encoding="utf-8")) + for name in ( + "lfm2_transformers_cuda.json", + "lfm2_transformers_cuda_zh.json", + "lfm2_transformers_cuda_long.json", + ) + ] + long_tokens = ( + references[0]["input_ids"] * 16 + ) # 208 tokens, spanning four 64-token pages. + cases = [ + ("A", references[0]["input_ids"]), + ("B", references[1]["input_ids"]), + ("C", references[2]["input_ids"]), + ("A", references[0]["input_ids"]), + ("long", long_tokens), + ("long", long_tokens), + ("long", long_tokens), + ] + capacity = max(len(tokens) for _, tokens in cases) + args.steps + 32 + block_size = 64 + num_blocks = max(8, (capacity + block_size - 1) // block_size) + cache = ( + _infinilm.StaticKVCacheConfig(1, capacity) + if args.cache_type == "static" + else _infinilm.PagedKVCacheConfig(num_blocks, block_size, 1) + ) + backend = "static-attn" if args.cache_type == "static" else "paged-attn" + + def build_input(engine, tokens, past): + total = past + len(tokens) + positions = list(range(past, total)) + kwargs = { + "mamba_init_state_indices": core.from_list( + [0 if past == 0 else 1], core.DataType.I32 + ), + "mamba_final_state_indices": core.from_list([1], core.DataType.I32), + } + if args.cache_type == "paged": + kwargs.update( + block_tables=core.from_list( + [list(range(num_blocks))], core.DataType.I32 + ), + slot_mapping=core.from_list(positions, core.DataType.I64), + ) + return engine._build_input( + core.from_list([tokens], core.DataType.I64), + position_ids=core.from_list([positions], core.DataType.I64), + past_kv_lengths=core.from_list([past], core.DataType.I32), + total_kv_lengths=core.from_list([total], core.DataType.I32), + input_offsets=core.from_list([0, len(tokens)], core.DataType.I32), + cu_seqlens=core.from_list([0, total], core.DataType.I32), + sample_all_positions=False, + **kwargs, + ) + + def float_logits(tensor): + if args.dtype == "float32": + array = tensor_to_numpy(core, tensor, np.float32) + else: + bits = tensor_to_numpy(core, tensor, np.uint16) + array = torch.from_numpy(bits).view(torch.bfloat16).float().numpy() + return array.reshape(-1, array.shape[-1])[-1] + + def gpu_memory(): + return int( + subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=memory.used", + "--format=csv,noheader,nounits", + ], + text=True, + ) + .strip() + .splitlines()[0] + ) + + runs, previous_sequences = [], {} + passed = True + with tempfile.TemporaryDirectory(prefix="lfm2-cache-validation-") as temporary: + config = json.loads((args.model / "config.json").read_text(encoding="utf-8")) + config["torch_dtype"] = config["dtype"] = args.dtype + (Path(temporary) / "config.json").write_text( + json.dumps(config), encoding="utf-8" + ) + engines = [] + try: + for _ in range(2 if args.compare_full_steps else 1): + engine = InferEngine( + temporary, + device=infinicore.device("cuda"), + cache_config=cache, + attention_backend=backend, + weight_load_mode="sync", + ) + load_model_state_dict_by_file( + engine, str(args.model), dtype=engine.dtype + ) + engines.append(engine) + cached_engine = engines[0] + full_engine = engines[1] if args.compare_full_steps else None + memory_after_load = gpu_memory() + for name, prompt_ids in cases: + history, generated, times, metrics = list(prompt_ids), [], [], [] + for step in range(args.steps): + chunk, past = ( + (history, 0) if step == 0 else ([history[-1]], len(history) - 1) + ) + native_input = build_input(cached_engine, chunk, past) + start = time.perf_counter() + output = _infinilm.InferEngine.forward(cached_engine, native_input) + core.sync_device() + times.append(time.perf_counter() - start) + actual = float_logits(output.logits) + if not np.isfinite(actual).all(): + raise AssertionError("Non-finite cached logits") + if step < args.compare_full_steps: + recomputed = _infinilm.InferEngine.forward( + full_engine, build_input(full_engine, history, 0) + ) + expected = float_logits(recomputed.logits) + if not np.isfinite(expected).all(): + raise AssertionError("Non-finite recomputed logits") + delta = actual - expected + close = bool( + np.allclose(actual, expected, atol=1e-4, rtol=1e-4) + ) + match = int(np.argmax(actual)) == int(np.argmax(expected)) + # Stable sorting preserves the minimum-ID tie convention. + candidates = sorted( + set( + np.argsort(-actual, kind="stable")[:5].tolist() + + np.argsort(-expected, kind="stable")[:5].tolist() + ) + ) + metrics.append( + { + "step": step, + "history_token_ids": list(history), + "cached_argmax": int(np.argmax(actual)), + "full_argmax": int(np.argmax(expected)), + "candidate_scores": [ + { + "token_id": int(candidate), + "cached": float(actual[candidate]), + "full": float(expected[candidate]), + } + for candidate in candidates + ], + "max_abs_error": float(np.abs(delta).max()), + "relative_l2_error": float( + np.linalg.norm(delta) + / max(np.linalg.norm(expected), 1e-12) + ), + "argmax_match": match, + "f32_tolerance_passed": close, + } + ) + if args.dtype == "float32" and not (close and match): + passed = False + token = int(np.argmax(actual)) + history.append(token) + generated.append(token) + consistent = ( + name not in previous_sequences + or generated == previous_sequences[name] + ) + passed = passed and consistent + previous_sequences[name] = generated + entry = { + "case": name, + "prompt_tokens": len(prompt_ids), + "generated_token_ids": generated, + "repeat_consistent": consistent, + "prefill_forward_ms": times[0] * 1000, + "decode_forward_mean_ms": float(np.mean(times[1:])) * 1000 + if len(times) > 1 + else None, + "decode_forward_tokens_per_second": (len(times) - 1) + / sum(times[1:]) + if len(times) > 1 + else None, + "gpu_memory_used_mib": gpu_memory(), + "full_cache_comparisons": metrics, + } + runs.append(entry) + print(json.dumps(entry), flush=True) + finally: + core.sync_device() + engines.clear() + result = { + "passed": passed, + "dtype": args.dtype, + "cache_type": args.cache_type, + "steps": args.steps, + "compare_full_steps": args.compare_full_steps, + "long_input_purpose": "synthetic repeated token history; cache stress, not text quality", + "timing_scope": "engine forward + device synchronization; excludes metadata, tokenization, transfer and loading", + "eos_policy": "ignore EOS for fixed-length kernel/cache stress", + "full_cache_numeric_gate": "atol=rtol=1e-4 plus argmax equality" + if args.dtype == "float32" + else "diagnostic only", + "model_engine_count": len(engines) + if engines + else (2 if args.compare_full_steps else 1), + "gpu_memory_scope": "nvidia-smi device-wide used MiB; not allocator peak", + "gpu_memory_after_load_mib": memory_after_load, + "runs": runs, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") + if not passed: + raise AssertionError("Cache consistency validation failed; see JSON") + + +if __name__ == "__main__": + main() From a3a30c556d208d35c575c789322e285df42ff281 Mon Sep 17 00:00:00 2001 From: xianyue <2869219499@qq.com> Date: Sat, 19 Sep 2026 20:14:32 +0800 Subject: [PATCH 2/2] docs: record clean CUDA regression --- LFM2_REPORT.md | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/LFM2_REPORT.md b/LFM2_REPORT.md index 4be464bb4..3e4c62d85 100644 --- a/LFM2_REPORT.md +++ b/LFM2_REPORT.md @@ -112,6 +112,7 @@ python test/models/lfm2/run_infinilm_lfm2_real.py \ --model /path/to/LFM2-1.2B \ --device cuda \ --cache-type static \ + --max-cache-len 256 \ --max-new-tokens 16 \ --repeat 2 \ --reference artifacts/lfm2_transformers_cuda.json \ @@ -120,12 +121,22 @@ python test/models/lfm2/run_infinilm_lfm2_real.py \ --output artifacts/lfm2_static_gate.json ``` -Paged Cache 使用相同命令，将 `--cache-type` 改为 `paged`。 +Paged Cache 使用相同命令，将 `--cache-type` 改为 `paged`，并将 `--max-cache-len` 改为 `1024`。 ## 4. 复现结果 ### 4.1 NVIDIA RTX 4090 +#### 2026-09-19 clean-build regression + +为排除旧构建缓存或预置扩展对结果的影响，在一台新创建的 RTX 4090 24 GB 实例上进行了从源码开始的复验。该实例使用 CUDA 12.8、Python 3.12.3、PyTorch `2.6.0a0+ecf3bae40a.nv25.01`。构建前仅将随源码传输带入的旧 `.xmake` 缓存和旧 `_infinilm` 扩展改名保留；随后完整重编译了 `_infinilm` 的全部 C++ 单元并安装。编译耗时 195.193 秒，新扩展 SHA-256 为 `9c7a361464aaaa82826a46cbf12cb89d5f57779317a4301068429c0ea3202945`。 + +- tiny F32 Full/Prefill+Decode：设置 `NVIDIA_TF32_OVERRIDE=0` 后，最大 logits 绝对误差为 `5.960464477539063e-08`，小于 `1e-5`；最终 argmax 一致。 +- Static KV Cache：`max_cache_len=256`，三种 Prompt 均与 Transformers 的 16 个 greedy token 精确一致；A/B/C/A/B/C 两轮结果一致。 +- Paged KV Cache：`max_cache_len=1024`，三种 Prompt 均与 Transformers 的 16 个 greedy token 精确一致；A/B/C/A/B/C 两轮结果一致。 + +本次结果 JSON 已归档到 `work/artifacts/final_4090_20260919/`：`lfm2_native_tiny_cuda_strict_new4090.json`、`lfm2_final_static_256_new4090.json` 和 `lfm2_final_paged_1024_new4090.json`。 + | 验收项 | Static | Paged | |---|---:|---:| | 三 Prompt Transformers 16-token 精确对齐 | 3/3 通过 | 3/3 通过 |