From 5a64b6ccaf0a9e10614338adae6c654003ee7a8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E6=81=92=E8=AF=AD=E6=AF=AB?= Date: Sun, 20 Sep 2026 23:49:20 +0800 Subject: [PATCH] feat(nvidia): add Mamba2 and RWKV5 adapters --- README.md | 65 +++ csrc/engine/compiler/general_compiler.cpp | 7 + csrc/models/infinilm_model.hpp | 1 + csrc/models/mamba2/mamba2_for_causal_lm.cpp | 347 ++++++++++++++ csrc/models/mamba2/mamba2_for_causal_lm.hpp | 95 ++++ csrc/models/rwkv5/rwkv5_for_causal_lm.cpp | 442 ++++++++++++++++++ csrc/models/rwkv5/rwkv5_for_causal_lm.hpp | 146 ++++++ python/infinilm/infer_engine.py | 26 +- python/infinilm/llm/llm.py | 2 + python/infinilm/llm/scheduler.py | 82 ++-- python/infinilm/modeling_utils.py | 105 +++++ python/infinilm/processors/__init__.py | 2 +- .../processors/basic_llm_processor.py | 15 +- .../infinilm/processors/mamba2_processor.py | 47 ++ python/infinilm/processors/rwkv5_processor.py | 180 +++++++ scripts/convert_rwkv5_checkpoint.py | 145 ++++++ scripts/prepare_mamba2_checkpoint.py | 66 +++ test/models/mamba2/REPORT.md | 76 +++ test/models/mamba2/__init__.py | 0 test/models/mamba2/benchmark.py | 128 +++++ test/models/mamba2/test_adaptation.py | 64 +++ test/models/mamba2/test_correctness.py | 244 ++++++++++ test/models/mamba2/test_real_model.py | 85 ++++ test/models/mamba2/test_scheduler.py | 43 ++ test/models/rwkv5/BENCHMARK.md | 85 ++++ test/models/rwkv5/benchmark.py | 205 ++++++++ test/models/rwkv5/test_adaptation.py | 135 ++++++ test/models/rwkv5/test_correctness.py | 345 ++++++++++++++ test/models/rwkv5/test_real_model.py | 82 ++++ test/models/rwkv5/test_scheduler.py | 87 ++++ xmake.lua | 12 +- 31 files changed, 3314 insertions(+), 50 deletions(-) create mode 100644 csrc/models/mamba2/mamba2_for_causal_lm.cpp create mode 100644 csrc/models/mamba2/mamba2_for_causal_lm.hpp create mode 100644 csrc/models/rwkv5/rwkv5_for_causal_lm.cpp create mode 100644 csrc/models/rwkv5/rwkv5_for_causal_lm.hpp create mode 100644 python/infinilm/processors/mamba2_processor.py create mode 100644 python/infinilm/processors/rwkv5_processor.py create mode 100644 scripts/convert_rwkv5_checkpoint.py create mode 100644 scripts/prepare_mamba2_checkpoint.py create mode 100644 test/models/mamba2/REPORT.md create mode 100644 test/models/mamba2/__init__.py create mode 100644 test/models/mamba2/benchmark.py create mode 100644 test/models/mamba2/test_adaptation.py create mode 100644 test/models/mamba2/test_correctness.py create mode 100644 test/models/mamba2/test_real_model.py create mode 100644 test/models/mamba2/test_scheduler.py create mode 100644 test/models/rwkv5/BENCHMARK.md create mode 100644 test/models/rwkv5/benchmark.py create mode 100644 test/models/rwkv5/test_adaptation.py create mode 100644 test/models/rwkv5/test_correctness.py create mode 100644 test/models/rwkv5/test_real_model.py create mode 100644 test/models/rwkv5/test_scheduler.py diff --git a/README.md b/README.md index 9a8b365f0..a96492dd2 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,71 @@ pip install -e . ``` + - RWKV-5（NVIDIA 单卡） + + 当前支持 RWKV-5.0、5.1 和 5.2 的官方 `.pth` 权重，以及 RWKV World + 字节词表。先将检查点转换为 InfiniLM 可加载的目录： + + ```bash + python scripts/convert_rwkv5_checkpoint.py \ + /models/RWKV-5-World-0.1B-v1-20230803-ctx4096.pth \ + /models/RWKV-5-World-0.1B-InfiniLM \ + --vocab /path/to/rwkv_vocab_v20230424.txt + ``` + + 在 NVIDIA GPU 上运行分页缓存推理： + + ```bash + python examples/test_infer.py \ + --device nvidia \ + --model /models/RWKV-5-World-0.1B-InfiniLM \ + --enable-paged-attn \ + --disable-prefix-caching \ + --num-blocks 64 \ + --block-size 16 \ + --max-new-tokens 32 + ``` + + RWKV 为每个请求保存独立的循环状态，不使用传统注意力 KV Cache。当前实现 + 限定 NVIDIA 单卡（TP=1、PP=1），并暂不支持前缀缓存、CUDA Graph、量化权重 + 和带可学习 `time_state` 的检查点。 + + RWKV5 的 paged 调度路径只分配循环状态行，因为 RWKV5 的架构没有注意力 KV + Cache。这里是模型架构适配，不是对通用 Transformer KV admission（包括未来 + token 预留策略）的优化；Transformer 和混合注意力模型仍按源代码使用 + `BlockManager`。可运行固定矩阵微基准： + + ```bash + INFINILM_RUN_GPU_TESTS=1 \ + RWKV5_MODEL_PATH=/models/RWKV-5-World-0.1B-InfiniLM \ + python test/models/rwkv5/benchmark.py \ + --cache-type paged --batch-sizes 1,2,4 \ + --input-lens 32,128,512 --output-lens 32,128 \ + --warmup 3 --runs 5 + ``` + + RTX 4090 D 的测试方法、两种调度路径对比和原始 JSONL 数据见 + `test/models/rwkv5/BENCHMARK.md`。 + + 静态缓存调度器目前是单请求串行模式，只使用 `--batch-sizes 1`。调度器专项单测 + 用于确认纯循环模型只占用循环状态行，同时确认带注意力的路径仍使用原来的 KV + block： + + ```bash + python test/models/rwkv5/test_scheduler.py -v + ``` + + 基础测试与 NVIDIA 正确性测试： + + ```bash + python test/models/rwkv5/test_adaptation.py -v + INFINILM_RUN_GPU_TESTS=1 \ + python test/models/rwkv5/test_correctness.py -v + INFINILM_RUN_GPU_TESTS=1 \ + RWKV5_MODEL_PATH=/models/RWKV-5-World-0.1B-InfiniLM \ + python test/models/rwkv5/test_real_model.py -v + ``` + - 单次推理测试 - llama示例 ```bash diff --git a/csrc/engine/compiler/general_compiler.cpp b/csrc/engine/compiler/general_compiler.cpp index 84ee670d4..d4d22366a 100644 --- a/csrc/engine/compiler/general_compiler.cpp +++ b/csrc/engine/compiler/general_compiler.cpp @@ -7,6 +7,9 @@ GeneralCompiler::GeneralCompiler(const std::shared_ptr &model, Ra } void GeneralCompiler::compile() { + if (!model_->supports_graph_compilation()) { + return; + } static_batching_compiler_->compile(); paged_compiler_->compile(); } @@ -14,6 +17,10 @@ void GeneralCompiler::compile() { GeneralCompiler::Compiled GeneralCompiler::get_compiled(const InfinilmModel::Input &input) { GeneralCompiler::Compiled result = {nullptr, nullptr}; + if (!model_->supports_graph_compilation()) { + return result; + } + // try each compiler, return the first valid result result = static_batching_compiler_.get()->get_compiled(input); if (std::get<0>(result) != nullptr && std::get<1>(result) != nullptr) { diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index 02677318d..f5ebce7d3 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -69,6 +69,7 @@ class InfinilmModel : public infinicore::nn::Module { virtual ~InfinilmModel() = default; virtual Output forward(const Input &input) const = 0; virtual void reset_cache(const cache::CacheConfig *cache_config); + virtual bool supports_graph_compilation() const { return true; } virtual const cache::CacheConfig *get_cache_config() const { return cache_config_.get(); } diff --git a/csrc/models/mamba2/mamba2_for_causal_lm.cpp b/csrc/models/mamba2/mamba2_for_causal_lm.cpp new file mode 100644 index 000000000..cc1179d66 --- /dev/null +++ b/csrc/models/mamba2/mamba2_for_causal_lm.cpp @@ -0,0 +1,347 @@ +#include "mamba2_for_causal_lm.hpp" + +#include "../../global_state/global_state.hpp" +#include "../models_registry.hpp" + +#include "infinicore/context/context.hpp" +#include "infinicore/ops/add.hpp" +#include "infinicore/ops/broadcast_to.hpp" +#include "infinicore/ops/causal_conv1d.hpp" +#include "infinicore/ops/mamba_selective_scan.hpp" +#include "infinicore/ops/mul.hpp" +#include "infinicore/ops/mul_scalar.hpp" +#include "infinicore/ops/silu.hpp" + +#include +#include +#include +#include + +namespace infinilm::models::mamba2 { +namespace { + +std::vector tensor_to_i32_vector(const infinicore::Tensor &tensor, + const char *name) { + if (!tensor || tensor->dtype() != infinicore::DataType::I32 || tensor->ndim() != 1) { + throw std::runtime_error(std::string("Mamba2: ") + name + + " must be a one-dimensional int32 tensor"); + } + auto cpu = tensor->device() == infinicore::Device::cpu() + ? tensor + : tensor->to(infinicore::Device::cpu()); + std::vector values(cpu->numel()); + std::memcpy(values.data(), cpu->data(), values.size() * sizeof(int32_t)); + return values; +} + +} // namespace + +std::shared_ptr +create_mamba2_model_config(std::shared_ptr config) { + if (config->get("model_type") != "mamba2") { + throw std::runtime_error("Mamba2 config creator called for a non-mamba2 model"); + } + auto &j = config->get_config_json(); + j["hidden_size"] = j.value("hidden_size", j.value("d_model", 768)); + j["num_hidden_layers"] = j.value("num_hidden_layers", j.value("n_layer", 24)); + j["intermediate_size"] = j.value( + "intermediate_size", j.value("d_inner", j["hidden_size"].get() * j.value("expand", 2))); + j["state_size"] = j.value("state_size", j.value("ssm_state_size", j.value("d_state", 64))); + j["conv_kernel"] = j.value("conv_kernel", j.value("d_conv", 4)); + j["num_heads"] = j.value("num_heads", j.value("nheads", 0)); + j["head_dim"] = j.value("head_dim", 0); + if (j["num_heads"].get() == 0 && j["head_dim"].get() == 0) { + j["head_dim"] = 64; + j["num_heads"] = j["intermediate_size"].get() / 64; + } else if (j["num_heads"].get() == 0) { + j["num_heads"] = j["intermediate_size"].get() / j["head_dim"].get(); + } else if (j["head_dim"].get() == 0) { + j["head_dim"] = j["intermediate_size"].get() / j["num_heads"].get(); + } + j["num_groups"] = j.value("num_groups", j.value("n_groups", 1)); + j["layer_norm_epsilon"] = j.value("layer_norm_epsilon", j.value("rms_norm_eps", 1e-5)); + j["rms_norm_eps"] = j.value("rms_norm_eps", j["layer_norm_epsilon"]); + j["norm_dim"] = j.value("norm_dim", j.value("mamba2_norm_dim", j["head_dim"])); + j["use_bias"] = j.value("use_bias", false); + j["use_conv_bias"] = j.value("use_conv_bias", true); + j["max_position_embeddings"] = j.value("max_position_embeddings", 8192); + + const size_t intermediate = j["intermediate_size"].get(); + const size_t heads = j["num_heads"].get(); + const size_t head_dim = j["head_dim"].get(); + const size_t groups = j["num_groups"].get(); + if (heads == 0 || head_dim == 0 || intermediate != heads * head_dim) { + throw std::runtime_error("Mamba2 requires intermediate_size == num_heads * head_dim"); + } + if (groups == 0 || heads % groups != 0) { + throw std::runtime_error("Mamba2 requires num_heads divisible by num_groups"); + } + if (groups != 1) { + throw std::runtime_error("Mamba2 adapter currently supports num_groups=1"); + } + return config; +} + +Mamba2Mixer::Mamba2Mixer(std::shared_ptr config, + size_t layer_idx, + const infinicore::Device &device) + : layer_idx_(layer_idx) { + const auto &dtype = config->get_dtype(); + hidden_size_ = config->get("hidden_size"); + intermediate_size_ = config->get("intermediate_size"); + state_size_ = config->get("state_size"); + num_heads_ = config->get("num_heads"); + head_dim_ = config->get("head_dim"); + const size_t norm_dim = config->get_or("norm_dim", head_dim_); + num_groups_ = config->get("num_groups"); + conv_kernel_ = config->get("conv_kernel"); + conv_dim_ = intermediate_size_ + 2 * num_groups_ * state_size_; + + const bool use_bias = config->get_or("use_bias", false); + const bool use_conv_bias = config->get_or("use_conv_bias", true); + const size_t projection_size = intermediate_size_ + conv_dim_ + num_heads_; + INFINICORE_NN_MODULE_INIT(in_proj, hidden_size_, projection_size, use_bias, dtype, device); + INFINICORE_NN_MODULE_INIT(out_proj, intermediate_size_, hidden_size_, use_bias, dtype, device); + INFINICORE_NN_MODULE_INIT(norm, norm_dim, config->get("rms_norm_eps"), dtype, device); + INFINICORE_NN_PARAMETER_INIT(conv1d_weight, ({conv_dim_, 1, conv_kernel_}, dtype, device)); + if (use_conv_bias) { + INFINICORE_NN_PARAMETER_INIT(conv1d_bias, ({conv_dim_}, dtype, device)); + } + INFINICORE_NN_PARAMETER_INIT(A_log, ({intermediate_size_}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(D, ({intermediate_size_}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(dt_bias, ({intermediate_size_}, dtype, device)); +} + +infinicore::Tensor Mamba2Mixer::forward( + const infinicore::Tensor &hidden_states, + const Mamba2BatchMetadata &metadata) const { + auto hidden_mut = const_cast(hidden_states); + auto projected = in_proj_->forward(hidden_mut); + auto z = projected->narrow({{2, 0, intermediate_size_}})->contiguous(); + auto xbc = projected->narrow({{2, intermediate_size_, conv_dim_}})->contiguous(); + auto dt = projected->narrow({{2, intermediate_size_ + conv_dim_, num_heads_}})->contiguous(); + + auto &context = infinilm::global_state::get_forward_context(); + if (layer_idx_ >= context.conv_state_vec.size() || !context.conv_state_vec[layer_idx_] + || layer_idx_ >= context.ssm_state_vec.size() || !context.ssm_state_vec[layer_idx_]) { + throw std::runtime_error("Mamba2 mixer state cache is not allocated"); + } + if (!context.mamba_metadata.input_offsets + || !context.mamba_metadata.init_state_indices + || !context.mamba_metadata.final_state_indices) { + throw std::runtime_error("Mamba2 mixer requires state metadata"); + } + auto conv_out = infinicore::op::causal_conv1d( + xbc, + context.conv_state_vec[layer_idx_], + conv1d_weight_, + conv1d_bias_ ? std::optional(conv1d_bias_) : std::nullopt, + context.mamba_metadata.input_offsets, + context.mamba_metadata.init_state_indices, + context.mamba_metadata.final_state_indices); + conv_out = infinicore::op::silu(conv_out); + auto x = conv_out->narrow({{2, 0, intermediate_size_}})->contiguous(); + auto b = conv_out->narrow({{2, intermediate_size_, state_size_}})->contiguous(); + auto c = conv_out->narrow({{2, intermediate_size_ + state_size_, state_size_}})->contiguous(); + + auto ssm_output = infinicore::Tensor::empty( + {hidden_states->size(0), hidden_states->size(1), intermediate_size_}, + hidden_states->dtype(), hidden_states->device()); + auto state_pool = context.ssm_state_vec[layer_idx_]; + const size_t request_count = metadata.input_offsets.size() - 1; + if (metadata.input_offsets.back() != static_cast(hidden_states->size(1))) { + throw std::runtime_error("Mamba2 input offsets do not cover hidden states"); + } + auto a_log_scan = infinicore::op::broadcast_to( + A_log_->view({intermediate_size_, 1}), + {intermediate_size_, state_size_}) + ->contiguous(); + auto d_scan = D_->view({intermediate_size_}); + auto dt_bias_scan = dt_bias_->view({intermediate_size_}); + + for (size_t request_idx = 0; request_idx < request_count; ++request_idx) { + const int32_t start = metadata.input_offsets[request_idx]; + const int32_t end = metadata.input_offsets[request_idx + 1]; + const int32_t read_index = metadata.init_state_indices[request_idx]; + const int32_t write_index = metadata.final_state_indices[request_idx]; + if (start < 0 || end <= start || read_index < 0 || write_index < 0 + || static_cast(read_index) >= state_pool->size(0) + || static_cast(write_index) >= state_pool->size(0)) { + throw std::runtime_error("Mamba2 received invalid state metadata"); + } + const size_t token_start = static_cast(start); + const size_t length = static_cast(end - start); + auto read_state = state_pool->narrow({{0, static_cast(read_index), 1}}); + infinicore::Tensor request_state; + if (read_index == write_index) { + request_state = read_state->view({1, intermediate_size_, state_size_}); + } else { + request_state = infinicore::Tensor::empty( + {1, intermediate_size_, state_size_}, infinicore::DataType::F32, + hidden_states->device()); + request_state->copy_from(read_state->view({1, intermediate_size_, state_size_})); + } + auto request_dt = dt->narrow({{1, token_start, length}})->view({1, length, num_heads_, 1}); + request_dt = infinicore::op::broadcast_to( + request_dt, {1, length, num_heads_, head_dim_}) + ->contiguous() + ->view({1, length, intermediate_size_}); + // The shared scan operator applies silu(gate) internally. Its neutral + // input is the positive solution of silu(x) = 1; Mamba2 applies its + // actual gate after per-head RMSNorm below. + auto neutral_gate = infinicore::op::mul_scalar( + infinicore::Tensor::ones( + x->narrow({{1, token_start, length}})->shape(), + x->dtype(), + x->device()), + 1.2784645427610737); + auto scan = infinicore::op::mamba_selective_scan( + x->narrow({{1, token_start, length}}), request_dt, + b->narrow({{1, token_start, length}}), + c->narrow({{1, token_start, length}}), + a_log_scan, d_scan, neutral_gate, + dt_bias_scan, request_state); + const size_t norm_dim = norm_->weight()->size(0); + infinicore::Tensor normalized; + if (norm_dim == intermediate_size_) { + normalized = norm_->forward(scan->view({length, intermediate_size_})); + normalized = normalized->view({1, length, intermediate_size_}); + } else if (norm_dim == head_dim_) { + normalized = norm_->forward(scan->view({length * num_heads_, head_dim_})); + normalized = normalized->view({1, length, intermediate_size_}); + } else { + throw std::runtime_error("Mamba2 norm_dim must equal intermediate_size or head_dim"); + } + normalized = infinicore::op::mul( + normalized, + infinicore::op::silu(z->narrow({{1, token_start, length}}))); + ssm_output->narrow({{1, token_start, length}})->copy_from(normalized); + if (read_index != write_index) { + state_pool->narrow({{0, static_cast(write_index), 1}}) + ->copy_from(request_state->view({1, num_heads_, head_dim_, state_size_})); + } + } + auto output_mut = ssm_output; + return out_proj_->forward(output_mut); +} + +Mamba2Block::Mamba2Block(std::shared_ptr config, + size_t layer_idx, + const infinicore::Device &device) { + const auto &dtype = config->get_dtype(); + const size_t hidden_size = config->get("hidden_size"); + INFINICORE_NN_MODULE_INIT(norm, hidden_size, config->get("rms_norm_eps"), dtype, device); + INFINICORE_NN_MODULE_INIT(mixer, config, layer_idx, device); +} + +infinicore::Tensor Mamba2Block::forward( + const infinicore::Tensor &hidden_states, + const Mamba2BatchMetadata &metadata) const { + auto normalized = norm_->forward(hidden_states); + return infinicore::op::add(hidden_states, mixer_->forward(normalized, metadata)); +} + +Mamba2Model::Mamba2Model(std::shared_ptr config, + const infinicore::Device &device) { + const auto &dtype = config->get_dtype(); + const size_t vocab_size = config->get("vocab_size"); + const size_t hidden_size = config->get("hidden_size"); + const size_t num_layers = config->get("num_hidden_layers"); + INFINICORE_NN_MODULE_INIT(embedding, vocab_size, hidden_size, std::nullopt, dtype, device); + layers_.reserve(num_layers); + for (size_t layer_idx = 0; layer_idx < num_layers; ++layer_idx) { + layers_.push_back(this->register_module( + "layers." + std::to_string(layer_idx), config, layer_idx, device)); + } + INFINICORE_NN_MODULE_INIT(norm_f, hidden_size, config->get("rms_norm_eps"), dtype, device); +} + +Mamba2BatchMetadata Mamba2Model::build_batch_metadata_( + const infinilm::InfinilmModel::Input &input) { + if (!input.input_offsets || !input.mamba_init_state_indices + || !input.mamba_final_state_indices) { + throw std::runtime_error("Mamba2 requires input offsets and state indices"); + } + Mamba2BatchMetadata metadata{ + tensor_to_i32_vector(*input.input_offsets, "input_offsets"), + tensor_to_i32_vector(*input.mamba_init_state_indices, "mamba_init_state_indices"), + tensor_to_i32_vector(*input.mamba_final_state_indices, "mamba_final_state_indices")}; + if (metadata.input_offsets.size() < 2 + || metadata.init_state_indices.size() + 1 != metadata.input_offsets.size() + || metadata.final_state_indices.size() != metadata.init_state_indices.size()) { + throw std::runtime_error("Mamba2 received inconsistent request metadata sizes"); + } + return metadata; +} + +infinicore::Tensor Mamba2Model::forward( + const infinilm::InfinilmModel::Input &input) const { + if (!input.input_ids) { + throw std::runtime_error("Mamba2 requires input_ids"); + } + const auto metadata = build_batch_metadata_(input); + auto input_ids = input.input_ids.value(); + if (input_ids->ndim() == 1) { + input_ids = input_ids->view({1, input_ids->size(0)}); + } + auto hidden_states = embedding_->forward(input_ids); + for (const auto &layer : layers_) { + hidden_states = layer->forward(hidden_states, metadata); + } + return norm_f_->forward(hidden_states); +} + +Mamba2ForCausalLM::Mamba2ForCausalLM( + std::shared_ptr config, + const infinicore::Device &device) + : infinilm::layers::causal_lm_templates::TextCausalLM( + std::move(config), device) {} + +void Mamba2ForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { + auto &context = infinilm::global_state::get_forward_context(); + context.kv_cache_vec.clear(); + context.conv_state_vec.clear(); + context.ssm_state_vec.clear(); + if (cache_config == nullptr) { + cache_config_.reset(); + return; + } + cache_config_ = cache_config->unique_copy(); + size_t pool_size = 0; + if (const auto *paged = dynamic_cast(cache_config)) { + pool_size = std::max(2, paged->num_blocks() / 4); + } else if (const auto *fixed = dynamic_cast(cache_config)) { + pool_size = fixed->max_batch_size() + 1; + } else { + throw std::runtime_error("Mamba2: unsupported cache configuration"); + } + const size_t num_layers = model_config_->get("num_hidden_layers"); + const size_t conv_dim = model_config_->get("intermediate_size") + + 2 * model_config_->get("num_groups") + * model_config_->get("state_size"); + const size_t conv_kernel = model_config_->get("conv_kernel"); + const size_t num_heads = model_config_->get("num_heads"); + const size_t head_dim = model_config_->get("head_dim"); + const size_t state_size = model_config_->get("state_size"); + const auto &dtype = model_config_->get_dtype(); + const auto device = infinicore::context::getDevice(); + context.conv_state_vec.reserve(num_layers); + context.ssm_state_vec.reserve(num_layers); + for (size_t layer_idx = 0; layer_idx < num_layers; ++layer_idx) { + context.conv_state_vec.push_back(infinicore::Tensor::zeros( + {pool_size, conv_dim, conv_kernel - 1}, dtype, device)); + context.ssm_state_vec.push_back(infinicore::Tensor::zeros( + {pool_size, num_heads, head_dim, state_size}, + infinicore::DataType::F32, device)); + } + infinicore::context::syncStream(); +} + +} // namespace infinilm::models::mamba2 + +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL( + mamba2, + infinilm::models::mamba2::Mamba2ForCausalLM, + infinilm::models::mamba2::create_mamba2_model_config); +} // namespace diff --git a/csrc/models/mamba2/mamba2_for_causal_lm.hpp b/csrc/models/mamba2/mamba2_for_causal_lm.hpp new file mode 100644 index 000000000..a149bb3ba --- /dev/null +++ b/csrc/models/mamba2/mamba2_for_causal_lm.hpp @@ -0,0 +1,95 @@ +#pragma once + +#include "../../cache/kv_cache.hpp" +#include "../../config/model_config.hpp" +#include "../../layers/causal_lm_templates/text_causal_lm.hpp" +#include "../../layers/linear/linear.hpp" +#include "../infinilm_model.hpp" + +#include "infinicore/nn/embedding.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/nn/rmsnorm.hpp" +#include "infinicore/tensor.hpp" + +#include +#include +#include +#include + +namespace infinilm::models::mamba2 { + +struct Mamba2BatchMetadata { + std::vector input_offsets; + std::vector init_state_indices; + std::vector final_state_indices; +}; + +class Mamba2Mixer : public infinicore::nn::Module { +public: + Mamba2Mixer(std::shared_ptr config, + size_t layer_idx, + const infinicore::Device &device); + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const Mamba2BatchMetadata &metadata) const; + +private: + size_t layer_idx_; + size_t hidden_size_; + size_t intermediate_size_; + size_t state_size_; + size_t num_heads_; + size_t head_dim_; + size_t num_groups_; + size_t conv_dim_; + size_t conv_kernel_; + + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, in_proj); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, out_proj); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); + INFINICORE_NN_PARAMETER(conv1d_weight); + INFINICORE_NN_PARAMETER(conv1d_bias); + INFINICORE_NN_PARAMETER(A_log); + INFINICORE_NN_PARAMETER(D); + INFINICORE_NN_PARAMETER(dt_bias); +}; + +class Mamba2Block : public infinicore::nn::Module { +public: + Mamba2Block(std::shared_ptr config, + size_t layer_idx, + const infinicore::Device &device); + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const Mamba2BatchMetadata &metadata) const; + +private: + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); + INFINICORE_NN_MODULE(Mamba2Mixer, mixer); +}; + +class Mamba2Model : public infinicore::nn::Module { +public: + Mamba2Model(std::shared_ptr config, + const infinicore::Device &device); + infinicore::Tensor forward(const infinilm::InfinilmModel::Input &input) const; + +private: + static Mamba2BatchMetadata build_batch_metadata_( + const infinilm::InfinilmModel::Input &input); + INFINICORE_NN_MODULE(infinicore::nn::Embedding, embedding); + INFINICORE_NN_MODULE_VEC(Mamba2Block, layers); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm_f); +}; + +class Mamba2ForCausalLM + : public infinilm::layers::causal_lm_templates::TextCausalLM { +public: + Mamba2ForCausalLM(std::shared_ptr config, + const infinicore::Device &device); + void reset_cache(const cache::CacheConfig *cache_config) override; + bool supports_graph_compilation() const override { return false; } +}; + +std::shared_ptr +create_mamba2_model_config(std::shared_ptr config); + +} // namespace infinilm::models::mamba2 diff --git a/csrc/models/rwkv5/rwkv5_for_causal_lm.cpp b/csrc/models/rwkv5/rwkv5_for_causal_lm.cpp new file mode 100644 index 000000000..92943c191 --- /dev/null +++ b/csrc/models/rwkv5/rwkv5_for_causal_lm.cpp @@ -0,0 +1,442 @@ +#include "rwkv5_for_causal_lm.hpp" + +#include "../../cache/cache.hpp" +#include "../../cache/kv_cache.hpp" +#include "../../global_state/global_state.hpp" +#include "../models_registry.hpp" + +#include "infinicore/context/context.hpp" +#include "infinicore/ops/add.hpp" +#include "infinicore/ops/layer_norm.hpp" +#include "infinicore/ops/lerp.hpp" +#include "infinicore/ops/mul.hpp" +#include "infinicore/ops/relu.hpp" +#include "infinicore/ops/rwkv5_wkv.hpp" +#include "infinicore/ops/sigmoid.hpp" +#include "infinicore/ops/silu.hpp" + +#include +#include +#include +#include + +namespace infinilm::models::rwkv5 { +namespace { + +std::vector tensor_to_i32_vector(const infinicore::Tensor &tensor, + const char *name) { + if (!tensor || tensor->dtype() != infinicore::DataType::I32 + || tensor->ndim() != 1) { + throw std::runtime_error(std::string("RWKV5: ") + name + + " must be a one-dimensional int32 tensor"); + } + auto cpu = tensor->device() == infinicore::Device::cpu() + ? tensor + : tensor->to(infinicore::Device::cpu()); + std::vector values(cpu->numel()); + std::memcpy(values.data(), cpu->data(), values.size() * sizeof(int32_t)); + return values; +} + +infinicore::Tensor time_mix(const infinicore::Tensor &previous, + const infinicore::Tensor ¤t, + const infinicore::Tensor &mix) { + auto broadcast_mix = mix->as_strided(current->shape(), {0, 0, 1}); + return infinicore::op::lerp(previous, current, broadcast_mix); +} + +infinicore::Tensor shift_with_state( + const infinicore::Tensor &hidden_states, + const infinicore::Tensor &state_pool, + const RWKV5BatchMetadata &metadata) { + const size_t hidden_size = hidden_states->size(2); + if (hidden_states->ndim() != 3 + || metadata.input_offsets.size() < 2 + || metadata.init_state_indices.size() + 1 + != metadata.input_offsets.size() + || metadata.final_state_indices.size() + != metadata.init_state_indices.size()) { + throw std::runtime_error("RWKV5: invalid shift-state tensor metadata"); + } + auto previous = infinicore::Tensor::empty( + hidden_states->shape(), hidden_states->dtype(), hidden_states->device()); + + const size_t request_count = metadata.input_offsets.size() - 1; + if (metadata.input_offsets.back() != static_cast(hidden_states->size(1))) { + throw std::runtime_error("RWKV5: shift-state offsets do not cover input"); + } + for (size_t request_idx = 0; request_idx < request_count; ++request_idx) { + const int32_t start = metadata.input_offsets[request_idx]; + const int32_t end = metadata.input_offsets[request_idx + 1]; + const int32_t read_index = metadata.init_state_indices[request_idx]; + const int32_t write_index = metadata.final_state_indices[request_idx]; + if (start < 0 || end <= start + || read_index < 0 || write_index < 0 + || static_cast(read_index) >= state_pool->size(0) + || static_cast(write_index) >= state_pool->size(0)) { + throw std::runtime_error("RWKV5: invalid shift-state metadata"); + } + + const size_t token_start = static_cast(start); + const size_t length = static_cast(end - start); + previous->narrow({{1, token_start, 1}})->copy_from( + state_pool->narrow( + {{0, static_cast(read_index), 1}}) + ->view({1, 1, hidden_size})); + if (length > 1) { + previous->narrow({{1, token_start + 1, length - 1}}) + ->copy_from(hidden_states->narrow({{1, token_start, length - 1}})); + } + state_pool->narrow({{0, static_cast(write_index), 1}}) + ->copy_from(hidden_states->narrow( + {{1, token_start + length - 1, 1}}) + ->view({1, hidden_size, 1})); + } + return previous; +} + +} // namespace + +std::shared_ptr +create_rwkv5_model_config(std::shared_ptr config) { + if (config->get("model_type") != "rwkv5") { + throw std::runtime_error("RWKV5 config creator called for a non-rwkv5 model"); + } + + auto &j = config->get_config_json(); + j["hidden_size"] = j.value("hidden_size", j.value("n_embd", 768)); + j["num_hidden_layers"] = j.value("num_hidden_layers", j.value("n_layer", 12)); + j["intermediate_size"] = j.value( + "intermediate_size", j.value("n_ffn", 4 * j["hidden_size"].get())); + j["num_attention_heads"] = j.value("num_attention_heads", j.value("n_head", 12)); + j["head_dim"] = j.value( + "head_dim", j["hidden_size"].get() / j["num_attention_heads"].get()); + j["layer_norm_eps"] = j.value("layer_norm_eps", 1e-5); + j["group_norm_eps"] = j.value("group_norm_eps", 64e-5); + j["use_attention_gate"] = j.value("use_attention_gate", true); + j["max_position_embeddings"] = j.value( + "max_position_embeddings", j.value("context_length", 4096)); + + const size_t hidden_size = j["hidden_size"].get(); + const size_t num_heads = j["num_attention_heads"].get(); + const size_t head_dim = j["head_dim"].get(); + if (num_heads == 0 || hidden_size != num_heads * head_dim) { + throw std::runtime_error( + "RWKV5 requires hidden_size == num_attention_heads * head_dim"); + } + return config; +} + +RWKV5HeadGroupNorm::RWKV5HeadGroupNorm(size_t hidden_size, + size_t num_heads, + double eps, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : num_heads_(num_heads), + head_size_(hidden_size / num_heads), + eps_(eps) { + INFINICORE_NN_PARAMETER_INIT(weight, ({hidden_size}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(bias, ({hidden_size}, dtype, device)); + unit_weight_ = infinicore::Tensor::ones({head_size_}, dtype, device); + zero_bias_ = infinicore::Tensor::zeros({head_size_}, dtype, device); +} + +infinicore::Tensor RWKV5HeadGroupNorm::forward( + const infinicore::Tensor &hidden_states) const { + const size_t token_count = hidden_states->numel() / (num_heads_ * head_size_); + auto by_head = hidden_states->view({token_count * num_heads_, head_size_}); + auto normalized = infinicore::op::layer_norm( + by_head, unit_weight_, zero_bias_, static_cast(eps_)); + normalized = normalized->view(hidden_states->shape()); + auto weight = weight_->as_strided(normalized->shape(), {0, 0, 1}); + auto bias = bias_->as_strided(normalized->shape(), {0, 0, 1}); + return infinicore::op::add(infinicore::op::mul(normalized, weight), bias); +} + +RWKV5TimeMix::RWKV5TimeMix( + std::shared_ptr config, + size_t layer_idx, + const infinicore::Device &device) + : layer_idx_(layer_idx) { + const auto &dtype = config->get_dtype(); + hidden_size_ = config->get("hidden_size"); + num_heads_ = config->get("num_attention_heads"); + head_size_ = config->get("head_dim"); + use_gate_ = config->get_or("use_attention_gate", true); + const double group_norm_eps = config->get("group_norm_eps"); + + INFINICORE_NN_PARAMETER_INIT(time_mix_k, ({hidden_size_}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(time_mix_v, ({hidden_size_}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(time_mix_r, ({hidden_size_}, dtype, device)); + if (use_gate_) { + INFINICORE_NN_PARAMETER_INIT(time_mix_g, ({hidden_size_}, dtype, device)); + } + INFINICORE_NN_PARAMETER_INIT(time_decay, ({num_heads_, head_size_}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(time_faaaa, ({num_heads_, head_size_}, dtype, device)); + + INFINICORE_NN_MODULE_INIT(key, hidden_size_, hidden_size_, false, dtype, device); + INFINICORE_NN_MODULE_INIT(value, hidden_size_, hidden_size_, false, dtype, device); + INFINICORE_NN_MODULE_INIT(receptance, hidden_size_, hidden_size_, false, dtype, device); + if (use_gate_) { + INFINICORE_NN_MODULE_INIT(gate, hidden_size_, hidden_size_, false, dtype, device); + } + INFINICORE_NN_MODULE_INIT(output, hidden_size_, hidden_size_, false, dtype, device); + INFINICORE_NN_MODULE_INIT( + ln_x, hidden_size_, num_heads_, group_norm_eps, dtype, device); +} + +infinicore::Tensor RWKV5TimeMix::run_wkv_( + const infinicore::Tensor &receptance, + const infinicore::Tensor &key_tensor, + const infinicore::Tensor &value_tensor, + const RWKV5BatchMetadata &metadata) const { + auto &states = infinilm::global_state::get_forward_context().ssm_state_vec; + if (layer_idx_ >= states.size() || !states[layer_idx_]) { + throw std::runtime_error("RWKV5TimeMix: WKV state cache is not allocated"); + } + auto state_pool = states[layer_idx_]; + auto out = infinicore::Tensor::empty( + receptance->shape(), receptance->dtype(), receptance->device()); + + const size_t request_count = metadata.input_offsets.size() - 1; + for (size_t request_idx = 0; request_idx < request_count; ++request_idx) { + const int32_t start = metadata.input_offsets[request_idx]; + const int32_t end = metadata.input_offsets[request_idx + 1]; + const int32_t read_index = metadata.init_state_indices[request_idx]; + const int32_t write_index = metadata.final_state_indices[request_idx]; + if (start < 0 || end <= start + || read_index < 0 || write_index < 0 + || static_cast(write_index) >= state_pool->size(0) + || static_cast(read_index) >= state_pool->size(0)) { + throw std::runtime_error("RWKV5TimeMix: invalid packed sequence metadata"); + } + + const size_t length = static_cast(end - start); + auto read_state = state_pool->narrow( + {{0, static_cast(read_index), 1}}); + infinicore::Tensor request_state; + if (read_index == write_index) { + request_state = read_state; + } else { + request_state = infinicore::Tensor::empty( + {1, num_heads_, head_size_, head_size_}, + infinicore::DataType::F32, + receptance->device()); + request_state->copy_from(read_state); + } + + infinicore::op::rwkv5_wkv_( + out->narrow({{1, static_cast(start), length}}), + receptance->narrow({{1, static_cast(start), length}}), + key_tensor->narrow({{1, static_cast(start), length}}), + value_tensor->narrow({{1, static_cast(start), length}}), + time_decay_, + time_faaaa_, + request_state); + + if (read_index != write_index) { + state_pool->narrow({{0, static_cast(write_index), 1}}) + ->copy_from(request_state); + } + } + return out; +} + +infinicore::Tensor RWKV5TimeMix::forward( + const infinicore::Tensor &hidden_states, + const RWKV5BatchMetadata &metadata) const { + auto &context = infinilm::global_state::get_forward_context(); + const size_t state_idx = layer_idx_ * 2; + if (state_idx >= context.conv_state_vec.size() + || !context.conv_state_vec[state_idx]) { + throw std::runtime_error("RWKV5TimeMix: time-mix state cache is not allocated"); + } + auto previous = shift_with_state( + hidden_states, context.conv_state_vec[state_idx], metadata); + + auto k_input = time_mix(previous, hidden_states, time_mix_k_); + auto v_input = time_mix(previous, hidden_states, time_mix_v_); + auto r_input = time_mix(previous, hidden_states, time_mix_r_); + auto k = key_->forward(k_input); + auto v = value_->forward(v_input); + auto r = receptance_->forward(r_input); + + auto mixed = ln_x_->forward(run_wkv_(r, k, v, metadata)); + if (use_gate_) { + auto g_input = time_mix(previous, hidden_states, time_mix_g_); + auto g = gate_->forward(g_input); + mixed = infinicore::op::mul(mixed, infinicore::op::silu(g)); + } + return output_->forward(mixed); +} + +RWKV5ChannelMix::RWKV5ChannelMix( + std::shared_ptr config, + size_t layer_idx, + const infinicore::Device &device) + : layer_idx_(layer_idx) { + const auto &dtype = config->get_dtype(); + const size_t hidden_size = config->get("hidden_size"); + const size_t intermediate_size = config->get("intermediate_size"); + INFINICORE_NN_PARAMETER_INIT(time_mix_k, ({hidden_size}, dtype, device)); + INFINICORE_NN_PARAMETER_INIT(time_mix_r, ({hidden_size}, dtype, device)); + INFINICORE_NN_MODULE_INIT(key, hidden_size, intermediate_size, false, dtype, device); + INFINICORE_NN_MODULE_INIT(value, intermediate_size, hidden_size, false, dtype, device); + INFINICORE_NN_MODULE_INIT(receptance, hidden_size, hidden_size, false, dtype, device); +} + +infinicore::Tensor RWKV5ChannelMix::forward( + const infinicore::Tensor &hidden_states, + const RWKV5BatchMetadata &metadata) const { + auto &context = infinilm::global_state::get_forward_context(); + const size_t state_idx = layer_idx_ * 2 + 1; + if (state_idx >= context.conv_state_vec.size() + || !context.conv_state_vec[state_idx]) { + throw std::runtime_error("RWKV5ChannelMix: channel-mix state cache is not allocated"); + } + auto previous = shift_with_state( + hidden_states, context.conv_state_vec[state_idx], metadata); + auto k_input = time_mix(previous, hidden_states, time_mix_k_); + auto r_input = time_mix(previous, hidden_states, time_mix_r_); + auto k = key_->forward(k_input); + k = infinicore::op::relu(k); + k = infinicore::op::mul(k, k); + auto value = value_->forward(k); + auto receptance = receptance_->forward(r_input); + return infinicore::op::mul(infinicore::op::sigmoid(receptance), value); +} + +RWKV5Block::RWKV5Block( + std::shared_ptr config, + size_t layer_idx, + const infinicore::Device &device) { + const auto &dtype = config->get_dtype(); + const size_t hidden_size = config->get("hidden_size"); + const double eps = config->get("layer_norm_eps"); + INFINICORE_NN_MODULE_INIT(ln1, hidden_size, eps, dtype, device); + INFINICORE_NN_MODULE_INIT(att, config, layer_idx, device); + INFINICORE_NN_MODULE_INIT(ln2, hidden_size, eps, dtype, device); + INFINICORE_NN_MODULE_INIT(ffn, config, layer_idx, device); +} + +infinicore::Tensor RWKV5Block::forward( + const infinicore::Tensor &hidden_states, + const RWKV5BatchMetadata &metadata) const { + auto x = ln1_->forward(hidden_states); + x = infinicore::op::add(hidden_states, att_->forward(x, metadata)); + auto channel_input = ln2_->forward(x); + return infinicore::op::add(x, ffn_->forward(channel_input, metadata)); +} + +RWKV5Model::RWKV5Model( + std::shared_ptr config, + const infinicore::Device &device) { + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + if (rank_info.tp_size != 1) { + throw std::runtime_error("RWKV5 currently supports tensor parallel size 1 only"); + } + const auto &dtype = config->get_dtype(); + const size_t vocab_size = config->get("vocab_size"); + const size_t hidden_size = config->get("hidden_size"); + const size_t num_layers = config->get("num_hidden_layers"); + const double eps = config->get("layer_norm_eps"); + INFINICORE_NN_MODULE_INIT( + embeddings, vocab_size, hidden_size, std::nullopt, dtype, device); + INFINICORE_NN_MODULE_INIT(ln0, hidden_size, eps, dtype, device); + blocks_.reserve(num_layers); + for (size_t layer_idx = 0; layer_idx < num_layers; ++layer_idx) { + blocks_.push_back(this->register_module( + "blocks." + std::to_string(layer_idx), config, layer_idx, device)); + } + INFINICORE_NN_MODULE_INIT(ln_out, hidden_size, eps, dtype, device); +} + +RWKV5BatchMetadata RWKV5Model::build_batch_metadata_( + const infinilm::InfinilmModel::Input &input) { + if (!input.input_offsets || !input.mamba_init_state_indices + || !input.mamba_final_state_indices) { + throw std::runtime_error( + "RWKV5 requires input offsets and initial/final state indices"); + } + RWKV5BatchMetadata metadata{ + tensor_to_i32_vector(*input.input_offsets, "input_offsets"), + tensor_to_i32_vector(*input.mamba_init_state_indices, "mamba_init_state_indices"), + tensor_to_i32_vector(*input.mamba_final_state_indices, "mamba_final_state_indices")}; + if (metadata.input_offsets.size() < 2 + || metadata.init_state_indices.size() + 1 != metadata.input_offsets.size() + || metadata.final_state_indices.size() != metadata.init_state_indices.size()) { + throw std::runtime_error("RWKV5 received inconsistent request metadata sizes"); + } + return metadata; +} + +infinicore::Tensor RWKV5Model::forward( + const infinilm::InfinilmModel::Input &input) const { + if (!input.input_ids) { + throw std::runtime_error("RWKV5 requires input_ids"); + } + const auto metadata = build_batch_metadata_(input); + auto hidden_states = ln0_->forward(embeddings_->forward(*input.input_ids)); + for (const auto &block : blocks_) { + hidden_states = block->forward(hidden_states, metadata); + } + return ln_out_->forward(hidden_states); +} + +RWKV5ForCausalLM::RWKV5ForCausalLM( + std::shared_ptr config, + const infinicore::Device &device) + : TextCausalLM(std::move(config), device) {} + +void RWKV5ForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { + auto &context = infinilm::global_state::get_forward_context(); + context.kv_cache_vec.clear(); + context.conv_state_vec.clear(); + context.ssm_state_vec.clear(); + if (cache_config == nullptr) { + cache_config_.reset(); + return; + } + + cache_config_ = cache_config->unique_copy(); + size_t pool_size = 0; + if (const auto *paged = dynamic_cast(cache_config)) { + pool_size = std::max(2, paged->num_blocks() / 4); + } else if (const auto *fixed = dynamic_cast(cache_config)) { + pool_size = fixed->max_batch_size() + 1; + } else { + throw std::runtime_error("RWKV5: unsupported cache configuration"); + } + + const size_t num_layers = model_config_->get("num_hidden_layers"); + const size_t hidden_size = model_config_->get("hidden_size"); + const size_t num_heads = model_config_->get("num_attention_heads"); + const size_t head_size = model_config_->get("head_dim"); + const auto &dtype = model_config_->get_dtype(); + const auto device = infinicore::context::getDevice(); + + context.conv_state_vec.reserve(num_layers * 2); + context.ssm_state_vec.reserve(num_layers); + for (size_t layer_idx = 0; layer_idx < num_layers; ++layer_idx) { + context.conv_state_vec.push_back( + infinicore::Tensor::zeros({pool_size, hidden_size, 1}, dtype, device)); + context.conv_state_vec.push_back( + infinicore::Tensor::zeros({pool_size, hidden_size, 1}, dtype, device)); + context.ssm_state_vec.push_back(infinicore::Tensor::zeros( + {pool_size, num_heads, head_size, head_size}, + infinicore::DataType::F32, + device)); + } + infinicore::context::syncStream(); +} + +} // namespace infinilm::models::rwkv5 + +namespace { + +INFINILM_REGISTER_CAUSAL_LM_MODEL( + rwkv5, + infinilm::models::rwkv5::RWKV5ForCausalLM, + infinilm::models::rwkv5::create_rwkv5_model_config); + +} // namespace diff --git a/csrc/models/rwkv5/rwkv5_for_causal_lm.hpp b/csrc/models/rwkv5/rwkv5_for_causal_lm.hpp new file mode 100644 index 000000000..bba68e549 --- /dev/null +++ b/csrc/models/rwkv5/rwkv5_for_causal_lm.hpp @@ -0,0 +1,146 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "../../layers/causal_lm_templates/text_causal_lm.hpp" +#include "../../layers/linear/linear.hpp" +#include "../infinilm_model.hpp" + +#include "infinicore/nn/embedding.hpp" +#include "infinicore/nn/layer_norm.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/nn/parameter.hpp" +#include "infinicore/tensor.hpp" + +#include +#include +#include + +namespace infinilm::models::rwkv5 { + +struct RWKV5BatchMetadata { + std::vector input_offsets; + std::vector init_state_indices; + std::vector final_state_indices; +}; + +class RWKV5HeadGroupNorm : public infinicore::nn::Module { +public: + RWKV5HeadGroupNorm(size_t hidden_size, + size_t num_heads, + double eps, + const infinicore::DataType &dtype, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +private: + size_t num_heads_; + size_t head_size_; + double eps_; + infinicore::Tensor unit_weight_; + infinicore::Tensor zero_bias_; + INFINICORE_NN_PARAMETER(weight); + INFINICORE_NN_PARAMETER(bias); +}; + +class RWKV5TimeMix : public infinicore::nn::Module { +public: + RWKV5TimeMix(std::shared_ptr config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const RWKV5BatchMetadata &metadata) const; + +private: + infinicore::Tensor run_wkv_(const infinicore::Tensor &receptance, + const infinicore::Tensor &key, + const infinicore::Tensor &value, + const RWKV5BatchMetadata &metadata) const; + + size_t layer_idx_; + size_t hidden_size_; + size_t num_heads_; + size_t head_size_; + bool use_gate_; + + INFINICORE_NN_PARAMETER(time_mix_k); + INFINICORE_NN_PARAMETER(time_mix_v); + INFINICORE_NN_PARAMETER(time_mix_r); + INFINICORE_NN_PARAMETER(time_mix_g); + INFINICORE_NN_PARAMETER(time_decay); + INFINICORE_NN_PARAMETER(time_faaaa); + + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, key); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, value); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, receptance); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, gate); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, output); + INFINICORE_NN_MODULE(RWKV5HeadGroupNorm, ln_x); +}; + +class RWKV5ChannelMix : public infinicore::nn::Module { +public: + RWKV5ChannelMix(std::shared_ptr config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const RWKV5BatchMetadata &metadata) const; + +private: + size_t layer_idx_; + INFINICORE_NN_PARAMETER(time_mix_k); + INFINICORE_NN_PARAMETER(time_mix_r); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, key); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, value); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, receptance); +}; + +class RWKV5Block : public infinicore::nn::Module { +public: + RWKV5Block(std::shared_ptr config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states, + const RWKV5BatchMetadata &metadata) const; + +private: + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, ln1); + INFINICORE_NN_MODULE(RWKV5TimeMix, att); + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, ln2); + INFINICORE_NN_MODULE(RWKV5ChannelMix, ffn); +}; + +class RWKV5Model : public infinicore::nn::Module { +public: + RWKV5Model(std::shared_ptr config, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinilm::InfinilmModel::Input &input) const; + +private: + static RWKV5BatchMetadata build_batch_metadata_( + const infinilm::InfinilmModel::Input &input); + + INFINICORE_NN_MODULE(infinicore::nn::Embedding, embeddings); + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, ln0); + INFINICORE_NN_MODULE_VEC(RWKV5Block, blocks); + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, ln_out); +}; + +class RWKV5ForCausalLM + : public infinilm::layers::causal_lm_templates::TextCausalLM { +public: + RWKV5ForCausalLM(std::shared_ptr config, + const infinicore::Device &device); + + void reset_cache(const cache::CacheConfig *cache_config) override; + bool supports_graph_compilation() const override { return false; } +}; + +std::shared_ptr +create_rwkv5_model_config(std::shared_ptr config); + +} // namespace infinilm::models::rwkv5 diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 117d82f9a..9d860779e 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -57,6 +57,10 @@ 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") == "mamba2" + or llm_config.get("model_type") == "mamba2" + or config.get("model_type") == "rwkv5" + or llm_config.get("model_type") == "rwkv5" or "linear_attention" in layer_types or all( key in llm_config @@ -242,11 +246,11 @@ def eos_token_id(self): # InternLM3's config.json has eos_token_id=2, while # generation_config.json has eos_token_id=[2, 128131]. # Following this priority ensures we always get the authoritative value. - eos_token_id = ( - self.hf_generation_config.get("eos_token_id") - or self.hf_config.get("eos_token_id") - or [] - ) + eos_token_id = self.hf_generation_config.get("eos_token_id") + if eos_token_id is None: + eos_token_id = self.hf_config.get("eos_token_id") + if eos_token_id is None: + eos_token_id = [] if isinstance(eos_token_id, int): eos_token_id = [eos_token_id] return eos_token_id @@ -452,6 +456,8 @@ def forward_raw( cu_seqlens=None, block_tables=None, slot_mapping=None, + mamba_init_state_indices=None, + mamba_final_state_indices=None, pixel_values=None, image_bound=None, tgt_sizes=None, @@ -474,6 +480,8 @@ def forward_raw( cu_seqlens=cu_seqlens, block_tables=block_tables, slot_mapping=slot_mapping, + mamba_init_state_indices=mamba_init_state_indices, + mamba_final_state_indices=mamba_final_state_indices, pixel_values=pixel_values, image_bound=image_bound, tgt_sizes=tgt_sizes, @@ -542,10 +550,10 @@ def generate( max_blocks_per_batch = 0 mamba_state_indices = None if self.has_mamba_cache and not self.enable_paged_attn: - if self.model_type != "mamba": - raise RuntimeError( - "Low-level generate for mamba-cache models currently requires paged attention" - ) + if self.model_type not in {"mamba", "mamba2", "rwkv5"}: + raise RuntimeError( + "Low-level generate for mamba-cache models currently requires paged attention" + ) elif self.has_mamba_cache: mamba_pool_size = max(2, self.get_cache_config().num_blocks() // 4) if batch_size > mamba_pool_size - 1: diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59d5a1eca..81ca42897 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -42,6 +42,7 @@ def __init__(self, config: EngineConfig): self.config = config hf_config = read_hf_config(config.model_path) has_mamba_cache = model_uses_mamba_cache(hf_config) + cacheless_state_model = hf_config.get("model_type") in {"mamba2", "rwkv5"} if has_mamba_cache and config.enable_prefix_caching: model_type = hf_config["model_type"] raise RuntimeError( @@ -108,6 +109,7 @@ def __init__(self, config: EngineConfig): max_num_batched_tokens=max_num_batched_tokens, connector=connector, has_mamba_cache=has_mamba_cache, + cacheless_state_model=cacheless_state_model, num_mamba_cache_blocks=num_mamba_cache_blocks, enable_prefix_caching=config.enable_prefix_caching, ) diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index c10b55f2f..1e655c21b 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -69,6 +69,7 @@ def __init__( max_num_batched_tokens: int = 1024, connector=None, has_mamba_cache: bool = False, + cacheless_state_model: bool = False, num_mamba_cache_blocks: int | None = None, enable_prefix_caching: bool = True, ): @@ -84,6 +85,7 @@ def __init__( self.cache_manager = BlockManager(num_blocks=num_blocks, block_size=block_size) self.has_mamba_cache = has_mamba_cache + self.cacheless_state_model = cacheless_state_model self.mamba_cache_manager = ( MambaCacheManager(num_mamba_cache_blocks or max(2, num_blocks // 4)) if has_mamba_cache @@ -199,37 +201,44 @@ def schedule(self) -> Optional[SchedulerOutput]: deferred_requests.append(req) break - if not self.can_accept_request( - req, - num_local_computed_tokens, - current_prefill_extra_blocks, - ): - logger.warning( - "Insufficient KV cache blocks for request %s, deferring.", - req.request_id, - ) - - if num_local_computed_tokens > 0: - self.cache_manager.free_blocks(cached_block_table) - deferred_requests.append(req) - break + # Pure recurrent models (Mamba2/RWKV5) do not have attention + # keys and values. Their state rows are managed separately; + # the normal BlockManager path remains unchanged for + # Transformer and hybrid attention models. + if self.cacheless_state_model: + req_blocks, slot_mapping = [], [] + else: + if not self.can_accept_request( + req, + num_local_computed_tokens, + current_prefill_extra_blocks, + ): + logger.warning( + "Insufficient KV cache blocks for request %s, deferring.", + req.request_id, + ) - allocation = self.cache_manager.allocate_slots( - num_new_tokens, - num_computed_tokens=num_computed_tokens, - cached_block_table=cached_block_table, - ) + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break - if allocation is None: - logger.warning( - "Failed to allocate KV cache blocks for request: %s", - req.request_id, + allocation = self.cache_manager.allocate_slots( + num_new_tokens, + num_computed_tokens=num_computed_tokens, + cached_block_table=cached_block_table, ) - if num_local_computed_tokens > 0: - self.cache_manager.free_blocks(cached_block_table) - deferred_requests.append(req) - break - req_blocks, slot_mapping = allocation + + if allocation is None: + logger.warning( + "Failed to allocate KV cache blocks for request: %s", + req.request_id, + ) + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break + req_blocks, slot_mapping = allocation if self.has_mamba_cache and req.mamba_cache_index is None: req.mamba_cache_index = self.mamba_cache_manager.allocate() @@ -317,11 +326,18 @@ def schedule(self) -> Optional[SchedulerOutput]: continue # Decode phase: allocate slot for newly generated token - req.block_table, new_slot = self.cache_manager.append_slot( - req.block_table, req.get_total_length() - ) - req.slot_mapping = [new_slot] - req.num_blocks = len(req.block_table) + # This is an architecture property, not a change to KV admission + # or the BlockManager reservation policy. + if self.cacheless_state_model: + req.block_table = [] + req.slot_mapping = [] + req.num_blocks = 0 + else: + req.block_table, new_slot = self.cache_manager.append_slot( + req.block_table, req.get_total_length() + ) + req.slot_mapping = [new_slot] + req.num_blocks = len(req.block_table) req.num_local_cached_tokens = req.get_total_length() - 1 scheduled_requests.append(req) diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index b210b0c56..7123e3b3c 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -741,6 +741,109 @@ def _remap_mamba(state_dict, config=None): return remapped +def _remap_mamba2(state_dict, config=None): + """Remap HuggingFace Mamba2 weights to the flattened scan layout.""" + config = config or {} + intermediate = int(config.get("intermediate_size", config.get("d_inner", 0))) + num_heads = int(config.get("num_heads", config.get("nheads", 0))) + head_dim = int(config.get("head_dim", 0)) + if intermediate == 0 and num_heads and head_dim: + intermediate = num_heads * head_dim + if num_heads == 0 and intermediate and head_dim: + num_heads = intermediate // head_dim + if head_dim == 0 and intermediate and num_heads: + head_dim = intermediate // num_heads + if not intermediate or not num_heads or not head_dim: + raise ValueError("Mamba2 config must define intermediate_size, num_heads and head_dim") + + remapped = {} + for key, tensor in state_dict.items(): + new_key = key.replace(".mixer.conv1d.weight", ".mixer.conv1d_weight") + new_key = new_key.replace(".mixer.conv1d.bias", ".mixer.conv1d_bias") + if new_key.startswith("backbone."): + new_key = "model." + new_key.removeprefix("backbone.") + if new_key.endswith((".mixer.A_log", ".mixer.D", ".mixer.dt_bias")): + values = tensor.reshape(-1) + if values.numel() == num_heads: + values = values.repeat_interleave(head_dim) + if values.numel() != intermediate: + raise ValueError( + f"Mamba2 parameter {key} has {values.numel()} values; " + f"expected {num_heads} or {intermediate}" + ) + tensor = values.contiguous() + remapped[new_key] = tensor + + if "lm_head.weight" not in remapped and "model.embedding.weight" in remapped: + remapped["lm_head.weight"] = remapped["model.embedding.weight"] + return remapped + + +def _remap_rwkv5(state_dict, config=None): + """Remap native RWKV-5 checkpoints to the InfiniLM module layout.""" + import torch + + config = config or {} + num_heads = int(config["num_attention_heads"]) + head_dim = int(config["head_dim"]) + remapped = {} + + def reshape_time_parameter(tensor, name, exponentiate=False): + original_dtype = tensor.dtype + values = tensor.squeeze() + if exponentiate: + values = torch.exp(values.float()).to(dtype=original_dtype) + if values.numel() == num_heads: + return ( + values.reshape(num_heads, 1) + .expand(num_heads, head_dim) + .contiguous() + ) + if values.numel() == num_heads * head_dim: + return values.reshape(num_heads, head_dim).contiguous() + raise ValueError( + f"RWKV-5 {name} has {values.numel()} values, expected " + f"{num_heads} or {num_heads * head_dim}" + ) + + for key, tensor in state_dict.items(): + if key.endswith(".time_state"): + raise ValueError( + "RWKV-5 checkpoints with learned time_state are not supported yet" + ) + + new_key = key + if key == "emb.weight": + new_key = "model.embeddings.weight" + elif key.startswith("blocks.0.ln0."): + new_key = "model.ln0." + key.removeprefix("blocks.0.ln0.") + elif key.startswith("blocks."): + new_key = "model." + key + elif key.startswith("ln_out."): + new_key = "model." + key + elif key == "head.weight": + new_key = "lm_head.weight" + + if new_key.endswith( + (".time_mix_k", ".time_mix_v", ".time_mix_r", ".time_mix_g") + ): + tensor = tensor.squeeze().reshape(-1).contiguous() + elif new_key.endswith(".time_decay"): + tensor = reshape_time_parameter(tensor, "time_decay") + elif new_key.endswith(".time_faaaa"): + tensor = reshape_time_parameter(tensor, "time_faaaa") + elif new_key.endswith(".time_first"): + # Older RWKV-5 checkpoints store log(time_first). + new_key = new_key.removesuffix(".time_first") + ".time_faaaa" + tensor = reshape_time_parameter( + tensor, "time_first", exponentiate=True + ) + + 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 +1180,8 @@ def _remap_kimi_k3(state_dict, config): "baichuan": _remap_baichuan, "gpt2": _remap_gpt2, "mamba": _remap_mamba, + "mamba2": _remap_mamba2, + "rwkv5": _remap_rwkv5, "videonsa": _remap_videonsa, "qwen3_5": _remap_qwen3_5, "ernie4_5_moe_vl": _remap_ernie4_5_moe_vl, diff --git a/python/infinilm/processors/__init__.py b/python/infinilm/processors/__init__.py index f1a543b69..75b3ca256 100644 --- a/python/infinilm/processors/__init__.py +++ b/python/infinilm/processors/__init__.py @@ -40,7 +40,7 @@ def from_pretrained(cls, model_dir_path: str, **kwargs) -> InfinilmProcessor: raw_config = json.load(f) raw_model_type = str(raw_config.get("model_type", "")).lower() - if raw_model_type in {"qwen3_5", "qwen3_5_moe"}: + if raw_model_type in {"qwen3_5", "qwen3_5_moe", "mamba2", "rwkv5"}: model_type = raw_model_type else: config = AutoConfig.from_pretrained(model_dir_path, trust_remote_code=True) diff --git a/python/infinilm/processors/basic_llm_processor.py b/python/infinilm/processors/basic_llm_processor.py index a6fbc33ac..d25c2018d 100644 --- a/python/infinilm/processors/basic_llm_processor.py +++ b/python/infinilm/processors/basic_llm_processor.py @@ -246,6 +246,17 @@ def _build_model_input_from_batch_scheduler_output( block_tables.append(padded_block_table) cu_seqlens.append(cu_seqlens[-1] + seq_len) + block_tables_tensor = ( + infinicore.from_list(block_tables, dtype=infinicore.int32) + if max_block_table_len > 0 + else None + ) + slot_mapping_tensor = ( + infinicore.from_list(slot_mapping, dtype=infinicore.int64) + if slot_mapping + else None + ) + return { "input_ids": infinicore.from_list([tokens], dtype=infinicore.int64), "position_ids": infinicore.from_list(position_ids, dtype=infinicore.int64), @@ -255,8 +266,8 @@ def _build_model_input_from_batch_scheduler_output( "total_kv_lengths": infinicore.from_list(seq_lens, dtype=infinicore.int32), "input_offsets": infinicore.from_list(seq_offsets, dtype=infinicore.int32), "cu_seqlens": infinicore.from_list(cu_seqlens, dtype=infinicore.int32), - "block_tables": infinicore.from_list(block_tables, dtype=infinicore.int32), - "slot_mapping": infinicore.from_list(slot_mapping, dtype=infinicore.int64), + "block_tables": block_tables_tensor, + "slot_mapping": slot_mapping_tensor, "temperature": temperature, "top_k": top_k, "top_p": top_p, diff --git a/python/infinilm/processors/mamba2_processor.py b/python/infinilm/processors/mamba2_processor.py new file mode 100644 index 000000000..24c5324bc --- /dev/null +++ b/python/infinilm/processors/mamba2_processor.py @@ -0,0 +1,47 @@ +import infinicore +from typing_extensions import override + +from ..llm.scheduler import SchedulerOutput +from ..llm.static_scheduler import StaticSchedulerOutput +from .mamba_processor import MambaProcessor +from .processor import register_processor + + +@register_processor("mamba2") +class Mamba2Processor(MambaProcessor): + """Processor for Mamba2 models with one recurrent state row per request.""" + + @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 + ) + init_indices = [] + final_indices = [] + for request in scheduler_output.scheduled_requests: + state_index = request.mamba_cache_index + if isinstance(scheduler_output, StaticSchedulerOutput): + state_index = 1 + if state_index is None: + raise RuntimeError( + f"Request {request.request_id} has no assigned Mamba2 state row" + ) + init_indices.append(0 if scheduler_output.is_prefill else state_index) + final_indices.append(state_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 + ) + model_inputs["block_tables"] = None + model_inputs["slot_mapping"] = None + return model_inputs diff --git a/python/infinilm/processors/rwkv5_processor.py b/python/infinilm/processors/rwkv5_processor.py new file mode 100644 index 000000000..4cba46570 --- /dev/null +++ b/python/infinilm/processors/rwkv5_processor.py @@ -0,0 +1,180 @@ +import ast +from pathlib import Path + +import infinicore +from typing_extensions import override + +from ..llm.scheduler import SchedulerOutput +from ..llm.static_scheduler import StaticSchedulerOutput +from .mamba_processor import MambaProcessor +from .processor import register_processor + + +class _TrieNode: + __slots__ = ("children", "token_id") + + def __init__(self): + self.children = {} + self.token_id = None + + +class RWKVWorldTokenizer: + """Greedy byte-trie tokenizer used by RWKV World checkpoints.""" + + def __init__(self, vocab_path: str | Path): + self.root = _TrieNode() + self.id_to_bytes = {} + with Path(vocab_path).open("r", encoding="utf-8") as vocab_file: + for line in vocab_file: + first_space = line.index(" ") + last_space = line.rindex(" ") + token_id = int(line[:first_space]) + token = ast.literal_eval(line[first_space:last_space]) + token_bytes = token.encode("utf-8") if isinstance(token, str) else token + expected_length = int(line[last_space:]) + if not isinstance(token_bytes, bytes) or len(token_bytes) != expected_length: + raise ValueError(f"Invalid RWKV vocabulary row for token {token_id}") + self.id_to_bytes[token_id] = token_bytes + node = self.root + for byte in token_bytes: + node = node.children.setdefault(byte, _TrieNode()) + node.token_id = token_id + + self.eos_token_id = 0 + self.bos_token_id = 0 + self.pad_token_id = 0 + self.vocab_size = len(self.id_to_bytes) + + def encode(self, text: str, add_special_tokens: bool = False, **kwargs): + del add_special_tokens, kwargs + source = text.encode("utf-8") + token_ids = [] + cursor = 0 + while cursor < len(source): + node = self.root + scan = cursor + best_end = None + best_id = None + while scan < len(source) and source[scan] in node.children: + node = node.children[source[scan]] + scan += 1 + if node.token_id is not None: + best_end = scan + best_id = node.token_id + if best_end is None: + raise ValueError(f"RWKV vocabulary cannot encode byte at offset {cursor}") + token_ids.append(best_id) + cursor = best_end + return token_ids + + def decode(self, token_ids, skip_special_tokens: bool = False, **kwargs): + del kwargs + pieces = [] + for token_id in token_ids: + token_id = int(token_id) + # RWKV reserves id 0 for EOS; it has no row in the byte vocabulary. + if token_id == self.eos_token_id: + continue + pieces.append(self.id_to_bytes[token_id]) + return b"".join(pieces).decode("utf-8", errors="replace") + + +@register_processor("rwkv5") +class RWKV5Processor(MambaProcessor): + def __init__(self, model_dir_path: str): + vocab_path = Path(model_dir_path) / "rwkv_vocab_v20230424.txt" + if not vocab_path.exists(): + raise FileNotFoundError( + f"RWKV tokenizer vocabulary not found: {vocab_path}" + ) + self.tokenizer = RWKVWorldTokenizer(vocab_path) + + @override + def __call__(self, prompt: str, return_tensors: str = None, **kwargs) -> dict: + del kwargs + token_ids = self.tokenizer.encode(prompt) + if return_tensors is None: + return {"input_ids": token_ids} + if return_tensors == "pt": + import torch + + return {"input_ids": torch.tensor([token_ids], dtype=torch.long)} + if return_tensors == "infini": + return { + "input_ids": infinicore.from_list( + [token_ids], dtype=infinicore.int64 + ) + } + raise ValueError(f"Unsupported return_tensors value: {return_tensors}") + + @override + def apply_chat_template( + self, + conversation, + add_generation_prompt: bool = False, + tokenize: bool = True, + **kwargs, + ): + del kwargs + sections = [] + for message in conversation: + content = message.get("content", "") + if isinstance(content, list): + content = "".join( + str(item.get("text", "")) if isinstance(item, dict) else str(item) + for item in content + ) + role = str(message.get("role", "user")).lower() + if role == "system": + sections.append(str(content).strip()) + elif role == "assistant": + sections.append(f"Assistant: {str(content).strip()}") + else: + sections.append(f"User: {str(content).strip()}") + if add_generation_prompt: + sections.append("Assistant:") + rendered = "\n\n".join(sections) + return self.tokenizer.encode(rendered) if tokenize else rendered + + @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, + ) + + init_indices = [] + final_indices = [] + for request in scheduler_output.scheduled_requests: + state_index = request.mamba_cache_index + if isinstance(scheduler_output, StaticSchedulerOutput): + state_index = 1 + if state_index is None: + raise RuntimeError( + f"Request {request.request_id} has no assigned RWKV state row" + ) + init_indices.append( + 0 if scheduler_output.is_prefill else state_index + ) + final_indices.append(state_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 + ) + # RWKV5 has no attention KV cache; only its recurrent state is needed. + model_inputs["block_tables"] = None + model_inputs["slot_mapping"] = None + return model_inputs diff --git a/scripts/convert_rwkv5_checkpoint.py b/scripts/convert_rwkv5_checkpoint.py new file mode 100644 index 000000000..2068410ba --- /dev/null +++ b/scripts/convert_rwkv5_checkpoint.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 +"""Convert an official RWKV-5 .pth checkpoint to InfiniLM safetensors.""" + +import argparse +import json +import shutil +from pathlib import Path + +import torch +from safetensors.torch import save_file + + +def _load_checkpoint(path: Path) -> dict[str, torch.Tensor]: + try: + checkpoint = torch.load(path, map_location="cpu", weights_only=True, mmap=True) + except TypeError: + checkpoint = torch.load(path, map_location="cpu") + if not isinstance(checkpoint, dict): + raise TypeError("RWKV checkpoint must contain a state dictionary") + return {key: value.contiguous() for key, value in checkpoint.items()} + + +def _infer_config(state_dict: dict[str, torch.Tensor], context_length: int) -> dict: + if any("time_maa" in key for key in state_dict): + raise ValueError("This converter supports RWKV-5 only, not RWKV-6 or newer") + if "blocks.0.att.time_decay" not in state_dict: + raise ValueError("Checkpoint does not contain RWKV-5 time_decay weights") + + embedding = state_dict["emb.weight"] + vocab_size, hidden_size = embedding.shape + layer_ids = { + int(key.split(".")[1]) + for key in state_dict + if key.startswith("blocks.") + } + num_layers = max(layer_ids) + 1 + intermediate_size = state_dict["blocks.0.ffn.key.weight"].shape[0] + time_decay = state_dict["blocks.0.att.time_decay"].squeeze() + if time_decay.ndim == 1: + # RWKV-5.0 stores one decay value per head. The value is shared by all + # channels in that head, so the converter expands it during loading. + num_heads = time_decay.numel() + if hidden_size % num_heads != 0: + raise ValueError( + f"hidden_size={hidden_size} is not divisible by heads={num_heads}" + ) + head_dim = hidden_size // num_heads + rwkv_version = "5.0" + elif time_decay.ndim == 2: + num_heads, head_dim = time_decay.shape + rwkv_version = "5.2" + else: + raise ValueError( + "RWKV-5 time_decay must have shape [num_heads] or " + "[num_heads, head_dim]" + ) + if hidden_size != num_heads * head_dim: + raise ValueError( + f"hidden_size={hidden_size} does not match heads={num_heads} x head_dim={head_dim}" + ) + + use_gate = "blocks.0.att.gate.weight" in state_dict + if rwkv_version == "5.0" and use_gate: + rwkv_version = "5.1" + if use_gate and "blocks.0.att.time_mix_g" not in state_dict: + raise ValueError("Gated RWKV-5 checkpoint is missing time_mix_g") + if not any( + key in state_dict + for key in ("blocks.0.att.time_first", "blocks.0.att.time_faaaa") + ): + raise ValueError("RWKV-5 checkpoint is missing time_first/time_faaaa") + if any(key.endswith(".time_state") for key in state_dict): + raise ValueError("Checkpoints with learned time_state are not supported yet") + + dtype_name = str(embedding.dtype).removeprefix("torch.") + return { + "architectures": ["RWKV5ForCausalLM"], + "model_type": "rwkv5", + "rwkv_version": rwkv_version, + "vocab_size": vocab_size, + "hidden_size": hidden_size, + "intermediate_size": intermediate_size, + "num_hidden_layers": num_layers, + "num_attention_heads": num_heads, + "head_dim": head_dim, + "context_length": context_length, + "max_position_embeddings": context_length, + "layer_norm_eps": 1e-5, + "group_norm_eps": 64e-5, + "use_attention_gate": use_gate, + "tie_word_embeddings": False, + "bos_token_id": 0, + "eos_token_id": 0, + "pad_token_id": 0, + "torch_dtype": dtype_name, + } + + +def convert(args: argparse.Namespace) -> None: + checkpoint_path = Path(args.checkpoint).expanduser().resolve() + output_dir = Path(args.output_dir).expanduser().resolve() + output_dir.mkdir(parents=True, exist_ok=True) + + state_dict = _load_checkpoint(checkpoint_path) + config = _infer_config(state_dict, args.context_length) + save_file( + state_dict, + output_dir / "model.safetensors", + metadata={"format": "pt", "source": checkpoint_path.name}, + ) + (output_dir / "config.json").write_text( + json.dumps(config, indent=2) + "\n", encoding="utf-8" + ) + (output_dir / "generation_config.json").write_text( + json.dumps( + {"bos_token_id": 0, "eos_token_id": 0, "pad_token_id": 0}, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + + vocab_path = Path(args.vocab).expanduser().resolve() + if not vocab_path.exists(): + raise FileNotFoundError(f"RWKV vocabulary does not exist: {vocab_path}") + shutil.copy2(vocab_path, output_dir / "rwkv_vocab_v20230424.txt") + print(f"Converted {checkpoint_path.name} -> {output_dir}") + print(json.dumps(config, indent=2)) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("checkpoint", help="Official RWKV-5 .pth checkpoint") + parser.add_argument("output_dir", help="Destination model directory") + parser.add_argument( + "--vocab", + required=True, + help="Path to rwkv_vocab_v20230424.txt", + ) + parser.add_argument("--context-length", type=int, default=4096) + return parser.parse_args() + + +if __name__ == "__main__": + convert(parse_args()) diff --git a/scripts/prepare_mamba2_checkpoint.py b/scripts/prepare_mamba2_checkpoint.py new file mode 100644 index 000000000..fb97af9c5 --- /dev/null +++ b/scripts/prepare_mamba2_checkpoint.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""Add InfiniLM model metadata to an official Mamba2 checkpoint directory.""" + +import argparse +import json +from pathlib import Path + + +def prepare_config(source: dict) -> dict: + """Translate the state-spaces config into the fields used by InfiniLM.""" + hidden_size = int(source.get("hidden_size", source.get("d_model", 768))) + intermediate_size = int( + source.get("intermediate_size", source.get("d_intermediate", 0)) or hidden_size * 2 + ) + # The released state-spaces Mamba2-130M config omits d_state; its + # projection and convolution shapes identify the trained value as 128. + state_size = int(source.get("state_size", source.get("d_state", 128))) + conv_kernel = int(source.get("conv_kernel", source.get("d_conv", 4))) + num_hidden_layers = int(source.get("num_hidden_layers", source.get("n_layer", 24))) + head_dim = int(source.get("head_dim", 64)) + num_heads = int(source.get("num_heads", intermediate_size // head_dim)) + + vocab_size = int(source.get("vocab_size", 0)) + vocab_multiple = int(source.get("pad_vocab_size_multiple", 1)) + if vocab_multiple > 1: + vocab_size = (vocab_size + vocab_multiple - 1) // vocab_multiple * vocab_multiple + + config = dict(source) + config.update( + { + "architectures": ["Mamba2ForCausalLM"], + "model_type": "mamba2", + "vocab_size": vocab_size, + "hidden_size": hidden_size, + "intermediate_size": intermediate_size, + "num_hidden_layers": num_hidden_layers, + "state_size": state_size, + "conv_kernel": conv_kernel, + "num_heads": num_heads, + "head_dim": head_dim, + "norm_dim": intermediate_size, + "num_groups": int(source.get("num_groups", 1)), + "rms_norm_eps": float(source.get("rms_norm_eps", 1e-5)), + "layer_norm_epsilon": float(source.get("layer_norm_epsilon", 1e-5)), + "use_bias": bool(source.get("use_bias", False)), + "use_conv_bias": bool(source.get("use_conv_bias", True)), + "torch_dtype": source.get("torch_dtype", "float16"), + "max_position_embeddings": int( + source.get("max_position_embeddings", source.get("max_seq_len", 2048)) + ), + } + ) + return config + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint_dir", type=Path) + args = parser.parse_args() + config_path = args.checkpoint_dir / "config.json" + config = json.loads(config_path.read_text()) + config_path.write_text(json.dumps(prepare_config(config), indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/test/models/mamba2/REPORT.md b/test/models/mamba2/REPORT.md new file mode 100644 index 000000000..a00c84c4f --- /dev/null +++ b/test/models/mamba2/REPORT.md @@ -0,0 +1,76 @@ +# Mamba2 NVIDIA Adapter Report + +## Scope + +- Model: `state-spaces/mamba2-130m` +- Platform: NVIDIA RTX 4090 D, 24 GB +- Runtime: InfiniLM + InfiniCore CUDA backend +- Supported path: single GPU, `num_groups=1`, paged scheduler with recurrent state rows +- Graph compilation: disabled for this adapter +- Tensor parallel: not implemented + +## Implementation + +The adapter implements the Mamba2 path in the existing InfiniLM model registry: + +1. Input projection is split into gate, SSM input, B/C parameters, and time-step values. +2. The existing CUDA causal-convolution operator updates one convolution history per request. +3. The existing CUDA selective-scan operator updates the recurrent SSM state. +4. RMSNorm and the SiLU gate are applied before the output projection. +5. Each request owns a state row, so requests do not share convolution or SSM history. +6. Mamba2 requests use only recurrent state rows; no Transformer attention KV + cache is needed by this pure recurrent architecture. + +The existing `BlockManager` and its KV admission calculations were left +unchanged. This adapter distinguishes a pure recurrent model from a hybrid +model that has both recurrent state and attention KV; it does not claim to fix +generic KV over-reservation. + +## Problems Found And Fixes + +### Build environment + +- The minimal remote xmake package had no Python module rule. The target is built as a shared library named `_infinilm.so` instead. +- The remote build needed the Python 3.12 include directory and `INFINI_ROOT=/data/InfiniCore/install`. +- Root xmake execution requires `xmake --root`. +- The first full rebuild used the default 130 parallel jobs and one compiler was killed by memory pressure. The successful rebuild used four jobs and one link job. + +### Test reference + +- The first PyTorch reference used the normalized tensor as the residual. The implementation correctly uses the original input as the residual. +- The reference did not write convolution and SSM states back to the request state table, hiding decode-state errors. Both state tables now persist after every forward. + +### Official checkpoint configuration + +- The released config omits `model_type`, Mamba2 dimensions, and the padded vocabulary size. A preparation script adds these fields. +- The embedding weight is `50288 x 768` although the original config says `50277`; the prepared config uses `50288`. +- The released projection and convolution shapes identify `state_size=128` for this checkpoint. The prepared config records that value. +- The official `mixer.norm.weight` has 1536 values, so the adapter supports a configurable full-intermediate RMSNorm in addition to the tiny-test head-dimension form. + +## Verification + +### Unit and correctness tests + +- Mamba2 checkpoint/config and weight-remapping tests: passed, 4 tests. +- Mamba2 scheduler state-cache test: passed, 1 test. +- Tiny random-weight GPU reference test: passed. It covers prefill, decode, recurrent state persistence, and request isolation with `atol=rtol=2e-4`. +- Real Mamba2-130M load/prefill/decode test: passed. + +### Real-model benchmark + +Input length is 128, generated length is 32, five measured runs after two warmups. + +| Batch | Prefill time | Decode time | Decode throughput | GPU used after benchmark | +| ---: | ---: | ---: | ---: | ---: | +| 1 | 145.4 ms | 158.3 ms | 202.1 tok/s | 1385 MiB | +| 4 | 577.1 ms | 322.5 ms | 396.9 tok/s | 1393 MiB | + +The observed GPU usage after loading was about 1307 MiB out of 24564 MiB. The value includes the process/runtime baseline reported by `nvidia-smi`, so it is not a pure parameter-size measurement. + +## Current Limitations + +- NVIDIA CUDA only; no Ascend or other accelerator implementation is included. +- Only `num_groups=1` is supported by the current scan integration. +- Tensor parallel and graph compilation are disabled for Mamba2. +- The official checkpoint repository does not include tokenizer files, so the verified real-model path uses raw token IDs. Text tokenizer and chat-generation verification need a compatible GPT-style tokenizer directory. +- The benchmark is a small raw-token benchmark, not a full service concurrency evaluation. diff --git a/test/models/mamba2/__init__.py b/test/models/mamba2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/models/mamba2/benchmark.py b/test/models/mamba2/benchmark.py new file mode 100644 index 000000000..358f91ca4 --- /dev/null +++ b/test/models/mamba2/benchmark.py @@ -0,0 +1,128 @@ +"""Small reproducible raw-token benchmark for the NVIDIA Mamba2 adapter.""" + +import argparse +import json +import os +import subprocess +import time + +import torch + +import infinicore +from infinilm.cache import PagedKVCacheConfig +from infinilm.distributed import DistConfig +from infinilm.infer_engine import InferEngine +from infinilm.modeling_utils import load_model_state_dict_by_file + + +def gpu_memory_mib(): + row = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=memory.used,memory.total", + "--format=csv,noheader,nounits", + ], + text=True, + ).strip() + used, total = (int(value.strip()) for value in row.split(",", 1)) + return {"used_mib": used, "total_mib": total} + + +def forward(engine, batch_size, input_len, step, state_rows): + decode = step > 0 + length = 1 if decode else input_len + tokens = [1] * (batch_size * length) + offsets = [index * length for index in range(batch_size + 1)] + past = [input_len + step - 1 if decode else 0] * batch_size + total = [value + length for value in past] + positions = [ + position + for past_length in past + for position in range(past_length, past_length + length) + ] + cu_seqlens = [0] + for value in total: + cu_seqlens.append(cu_seqlens[-1] + value) + return engine.forward_raw( + infinicore.from_list([tokens], dtype=infinicore.int64), + position_ids=infinicore.from_list(positions, dtype=infinicore.int64), + past_kv_lengths=infinicore.from_list(past, dtype=infinicore.int32), + total_kv_lengths=infinicore.from_list(total, dtype=infinicore.int32), + input_offsets=infinicore.from_list(offsets, dtype=infinicore.int32), + cu_seqlens=infinicore.from_list(cu_seqlens, dtype=infinicore.int32), + mamba_init_state_indices=infinicore.from_list( + [0 if not decode else row for row in state_rows], + dtype=infinicore.int32, + ), + mamba_final_state_indices=infinicore.from_list( + state_rows, dtype=infinicore.int32 + ), + sample_all_positions=False, + ) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", default=os.environ.get("MAMBA2_MODEL_PATH"), required=True) + parser.add_argument("--batch-size", type=int, default=4) + parser.add_argument("--input-len", type=int, default=128) + parser.add_argument("--output-len", type=int, default=32) + parser.add_argument("--runs", type=int, default=5) + args = parser.parse_args() + + engine = InferEngine( + args.model, + device=infinicore.device("cuda", 0), + distributed_config=DistConfig(1), + cache_config=PagedKVCacheConfig( + num_blocks=64, block_size=16, max_batch_size=args.batch_size + ), + attention_backend="paged-attn", + weight_load_mode="sync", + ) + load_model_state_dict_by_file(engine, args.model, dtype=infinicore.float16) + rows = list(range(1, args.batch_size + 1)) + print(json.dumps({"memory_after_load": gpu_memory_mib()}), flush=True) + + for _ in range(2): + forward(engine, args.batch_size, args.input_len, 0, rows) + for step in range(1, args.output_len + 1): + forward(engine, args.batch_size, args.input_len, step, rows) + torch.cuda.synchronize() + + prefill_samples = [] + decode_samples = [] + for _ in range(args.runs): + start = time.perf_counter() + forward(engine, args.batch_size, args.input_len, 0, rows) + torch.cuda.synchronize() + prefill_samples.append(time.perf_counter() - start) + + start = time.perf_counter() + for step in range(1, args.output_len + 1): + forward(engine, args.batch_size, args.input_len, step, rows) + torch.cuda.synchronize() + decode_samples.append(time.perf_counter() - start) + + prefill = sum(prefill_samples) / len(prefill_samples) + decode = sum(decode_samples) / len(decode_samples) + print( + json.dumps( + { + "batch_size": args.batch_size, + "input_len": args.input_len, + "output_len": args.output_len, + "prefill_ms": prefill * 1000, + "decode_ms": decode * 1000, + "decode_tokens_per_second": ( + args.batch_size * args.output_len / decode + ), + "memory_after_benchmark": gpu_memory_mib(), + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/models/mamba2/test_adaptation.py b/test/models/mamba2/test_adaptation.py new file mode 100644 index 000000000..42600de39 --- /dev/null +++ b/test/models/mamba2/test_adaptation.py @@ -0,0 +1,64 @@ +import json +import unittest + +import torch + +from infinilm.infer_engine import model_uses_mamba_cache +from infinilm.modeling_utils import _remap_mamba2 +from scripts.prepare_mamba2_checkpoint import prepare_config + + +class Mamba2CheckpointTest(unittest.TestCase): + def test_prepares_official_state_spaces_config(self): + config = prepare_config( + { + "d_model": 768, + "d_intermediate": 0, + "n_layer": 24, + "vocab_size": 50277, + "pad_vocab_size_multiple": 16, + "ssm_cfg": {"layer": "Mamba2"}, + } + ) + + self.assertEqual(config["model_type"], "mamba2") + self.assertEqual(config["vocab_size"], 50288) + self.assertEqual(config["hidden_size"], 768) + self.assertEqual(config["intermediate_size"], 1536) + self.assertEqual(config["state_size"], 128) + self.assertEqual(config["num_heads"], 24) + self.assertEqual(config["head_dim"], 64) + + def test_remaps_official_prefixes_and_expands_head_parameters(self): + state_dict = { + "backbone.embedding.weight": torch.zeros(32, 8), + "backbone.layers.0.mixer.A_log": torch.tensor([1.0, 2.0]), + "backbone.layers.0.mixer.D": torch.tensor([3.0, 4.0]), + "backbone.layers.0.mixer.dt_bias": torch.tensor([5.0, 6.0]), + } + remapped = _remap_mamba2( + state_dict, + {"intermediate_size": 8, "num_heads": 2, "head_dim": 4}, + ) + + self.assertIn("model.embedding.weight", remapped) + self.assertIn("lm_head.weight", remapped) + self.assertEqual(remapped["model.layers.0.mixer.A_log"].shape, (8,)) + torch.testing.assert_close( + remapped["model.layers.0.mixer.A_log"], + torch.tensor([1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0]), + ) + + def test_rejects_inconsistent_parameter_size(self): + with self.assertRaisesRegex(ValueError, "expected 2 or 8"): + _remap_mamba2( + {"backbone.layers.0.mixer.A_log": torch.zeros(3)}, + {"intermediate_size": 8, "num_heads": 2, "head_dim": 4}, + ) + + def test_model_uses_recurrent_state_cache(self): + self.assertTrue(model_uses_mamba_cache({"model_type": "mamba2"})) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/mamba2/test_correctness.py b/test/models/mamba2/test_correctness.py new file mode 100644 index 000000000..f77a9a4e4 --- /dev/null +++ b/test/models/mamba2/test_correctness.py @@ -0,0 +1,244 @@ +import json +import os +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from safetensors.torch import save_file + +import infinicore +from infinilm.cache import PagedKVCacheConfig +from infinilm.distributed import DistConfig +from infinilm.infer_engine import InferEngine +from infinilm.modeling_utils import load_model_state_dict_by_file + + +RUN_GPU_TESTS = os.getenv("INFINILM_RUN_GPU_TESTS") == "1" + + +def make_tiny_checkpoint(seed=20260920): + generator = torch.Generator().manual_seed(seed) + config = { + "architectures": ["Mamba2ForCausalLM"], + "model_type": "mamba2", + "vocab_size": 32, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 1, + "state_size": 4, + "conv_kernel": 4, + "num_heads": 2, + "head_dim": 8, + "num_groups": 1, + "rms_norm_eps": 1e-5, + "use_bias": False, + "use_conv_bias": True, + "torch_dtype": "float32", + "max_position_embeddings": 128, + "bos_token_id": 0, + "eos_token_id": 0, + "pad_token_id": 0, + } + vocab = config["vocab_size"] + hidden = config["hidden_size"] + inner = config["intermediate_size"] + heads = config["num_heads"] + state_size = config["state_size"] + conv_dim = inner + 2 * state_size + projection_size = inner + conv_dim + heads + + def randn(*shape, scale=0.08): + return torch.randn(*shape, generator=generator) * scale + + mixer = "backbone.layers.0.mixer" + weights = { + "backbone.embedding.weight": randn(vocab, hidden), + "backbone.norm_f.weight": torch.ones(hidden), + "lm_head.weight": randn(vocab, hidden), + "backbone.layers.0.norm.weight": torch.ones(hidden), + f"{mixer}.norm.weight": torch.ones(config["head_dim"]), + f"{mixer}.in_proj.weight": randn(projection_size, hidden), + f"{mixer}.out_proj.weight": randn(hidden, inner), + f"{mixer}.conv1d.weight": randn(conv_dim, 1, config["conv_kernel"]), + f"{mixer}.conv1d.bias": randn(conv_dim, scale=0.01), + f"{mixer}.A_log": torch.zeros(heads), + f"{mixer}.D": torch.ones(heads), + f"{mixer}.dt_bias": torch.zeros(heads), + } + return weights, config + + +class Mamba2Reference: + def __init__(self, weights, config): + self.w = weights + self.config = config + self.conv_states = {} + self.ssm_states = {} + + def _state(self, state_id): + inner = self.config["intermediate_size"] + dstate = self.config["state_size"] + heads = self.config["num_heads"] + head_dim = self.config["head_dim"] + conv_dim = inner + 2 * dstate + self.conv_states.setdefault( + state_id, + torch.zeros(conv_dim, self.config["conv_kernel"] - 1), + ) + self.ssm_states.setdefault( + state_id, + torch.zeros(heads, head_dim, dstate), + ) + return self.conv_states[state_id], self.ssm_states[state_id] + + def _rms_norm(self, x, weight): + eps = self.config["rms_norm_eps"] + return x * torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + eps) * weight + + def forward(self, token_ids, state_id): + c = self.config + mixer = "backbone.layers.0.mixer" + conv_state, ssm_state = self._state(state_id) + x = F.embedding(torch.tensor(token_ids), self.w["backbone.embedding.weight"]) + residual = x + x = self._rms_norm(x, self.w["backbone.layers.0.norm.weight"]) + + projected = F.linear(x, self.w[f"{mixer}.in_proj.weight"]) + inner = c["intermediate_size"] + dstate = c["state_size"] + heads = c["num_heads"] + head_dim = c["head_dim"] + z = projected[:, :inner] + xbc = projected[:, inner : inner + inner + 2 * dstate] + dt = projected[:, inner + inner + 2 * dstate :] + + conv_weight = self.w[f"{mixer}.conv1d.weight"].squeeze(1) + conv_bias = self.w[f"{mixer}.conv1d.bias"] + conv_outputs = [] + for token in xbc: + window = torch.cat((conv_state, token[:, None]), dim=1) + conv_outputs.append((window * conv_weight).sum(dim=1) + conv_bias) + conv_state = window[:, 1:].clone() + conv_outputs = F.silu(torch.stack(conv_outputs)) + x_part = conv_outputs[:, :inner].view(-1, heads, head_dim) + b_part = conv_outputs[:, inner : inner + dstate] + c_part = conv_outputs[:, inner + dstate :] + + ys = [] + a = -torch.exp(self.w[f"{mixer}.A_log"]) + d = self.w[f"{mixer}.D"] + dt_bias = self.w[f"{mixer}.dt_bias"] + for token_idx in range(len(token_ids)): + dt_token = F.softplus(dt[token_idx] + dt_bias) + ssm_state = ( + torch.exp(dt_token[:, None, None] * a[:, None, None]) * ssm_state + + dt_token[:, None, None] + * x_part[token_idx][:, :, None] + * b_part[token_idx][None, None, :] + ) + y = (ssm_state * c_part[token_idx][None, None, :]).sum(dim=-1) + y = y + x_part[token_idx] * d[:, None] + y = self._rms_norm(y, self.w[f"{mixer}.norm.weight"]) + ys.append((y * F.silu(z[token_idx].view(heads, head_dim))).reshape(inner)) + + mixed = F.linear(torch.stack(ys), self.w[f"{mixer}.out_proj.weight"]) + self.conv_states[state_id] = conv_state + self.ssm_states[state_id] = ssm_state + hidden = residual + mixed + hidden = self._rms_norm(hidden, self.w["backbone.norm_f.weight"]) + return F.linear(hidden, self.w["lm_head.weight"]) + + +@unittest.skipUnless( + RUN_GPU_TESTS and torch.cuda.is_available(), + "set INFINILM_RUN_GPU_TESTS=1 on an NVIDIA host", +) +class Mamba2GPUCorrectnessTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.temp_dir = tempfile.TemporaryDirectory() + model_dir = Path(cls.temp_dir.name) + cls.weights, cls.config = make_tiny_checkpoint() + save_file(cls.weights, model_dir / "model.safetensors", metadata={"format": "pt"}) + (model_dir / "config.json").write_text(json.dumps(cls.config)) + cls.engine = InferEngine( + str(model_dir), + device=infinicore.device("cuda", 0), + distributed_config=DistConfig(1), + cache_config=PagedKVCacheConfig( + num_blocks=16, block_size=16, max_batch_size=4 + ), + attention_backend="paged-attn", + weight_load_mode="sync", + ) + load_model_state_dict_by_file( + cls.engine, str(model_dir), dtype=infinicore.float32 + ) + + @classmethod + def tearDownClass(cls): + del cls.engine + cls.temp_dir.cleanup() + + def forward(self, requests, initial_rows, final_rows, past_lengths): + flat_tokens = [token for request in requests for token in request] + offsets = [0] + for request in requests: + offsets.append(offsets[-1] + len(request)) + total_lengths = [ + past + len(request) for past, request in zip(past_lengths, requests) + ] + positions = [ + position + for past, request in zip(past_lengths, requests) + for position in range(past, past + len(request)) + ] + cu_seqlens = [0] + for total_length in total_lengths: + cu_seqlens.append(cu_seqlens[-1] + total_length) + output = self.engine.forward_raw( + infinicore.from_list([flat_tokens], dtype=infinicore.int64), + position_ids=infinicore.from_list(positions, dtype=infinicore.int64), + past_kv_lengths=infinicore.from_list(past_lengths, dtype=infinicore.int32), + total_kv_lengths=infinicore.from_list(total_lengths, dtype=infinicore.int32), + input_offsets=infinicore.from_list(offsets, dtype=infinicore.int32), + cu_seqlens=infinicore.from_list(cu_seqlens, dtype=infinicore.int32), + mamba_init_state_indices=infinicore.from_list( + initial_rows, dtype=infinicore.int32 + ), + mamba_final_state_indices=infinicore.from_list( + final_rows, dtype=infinicore.int32 + ), + sample_all_positions=True, + ) + return torch.from_numpy(np.asarray(output["logits"].to_numpy())).reshape( + len(flat_tokens), -1 + ) + + def test_prefill_decode_and_state_isolation(self): + reference = Mamba2Reference(self.weights, self.config) + + actual = self.forward([[1, 2, 3]], [0], [1], [0]) + expected = reference.forward([1, 2, 3], 1) + torch.testing.assert_close(actual, expected, atol=2e-4, rtol=2e-4) + + actual = self.forward([[4]], [1], [1], [3]) + expected = reference.forward([4], 1) + torch.testing.assert_close(actual, expected, atol=2e-4, rtol=2e-4) + + actual = self.forward([[5, 6], [7, 8]], [0, 0], [2, 3], [0, 0]) + expected = torch.cat( + [reference.forward([5, 6], 2), reference.forward([7, 8], 3)] + ) + torch.testing.assert_close(actual, expected, atol=2e-4, rtol=2e-4) + + actual = self.forward([[9], [10]], [2, 3], [2, 3], [2, 2]) + expected = torch.cat([reference.forward([9], 2), reference.forward([10], 3)]) + torch.testing.assert_close(actual, expected, atol=2e-4, rtol=2e-4) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/mamba2/test_real_model.py b/test/models/mamba2/test_real_model.py new file mode 100644 index 000000000..a0cb6c35f --- /dev/null +++ b/test/models/mamba2/test_real_model.py @@ -0,0 +1,85 @@ +import ctypes +import os +import unittest + +import numpy as np +import torch + +import infinicore +from infinilm.cache import PagedKVCacheConfig +from infinilm.distributed import DistConfig +from infinilm.infer_engine import InferEngine +from infinilm.modeling_utils import load_model_state_dict_by_file + + +MODEL_PATH = os.getenv("MAMBA2_MODEL_PATH") +RUN_GPU_TESTS = os.getenv("INFINILM_RUN_GPU_TESTS") == "1" + + +@unittest.skipUnless( + RUN_GPU_TESTS and MODEL_PATH and torch.cuda.is_available(), + "set INFINILM_RUN_GPU_TESTS=1 and MAMBA2_MODEL_PATH on an NVIDIA host", +) +class Mamba2RealModelTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.engine = InferEngine( + MODEL_PATH, + device=infinicore.device("cuda", 0), + distributed_config=DistConfig(1), + cache_config=PagedKVCacheConfig( + num_blocks=64, block_size=16, max_batch_size=4 + ), + attention_backend="paged-attn", + weight_load_mode="sync", + ) + load_model_state_dict_by_file(cls.engine, MODEL_PATH, dtype=infinicore.float16) + + @classmethod + def tearDownClass(cls): + del cls.engine + + def _forward(self, tokens, initial_row, final_row, past_length): + length = len(tokens) + output = self.engine.forward_raw( + infinicore.from_list([tokens], dtype=infinicore.int64), + position_ids=infinicore.from_list( + [list(range(past_length, past_length + length))], + dtype=infinicore.int64, + ), + past_kv_lengths=infinicore.from_list([past_length], dtype=infinicore.int32), + total_kv_lengths=infinicore.from_list( + [past_length + length], dtype=infinicore.int32 + ), + input_offsets=infinicore.from_list([0, length], dtype=infinicore.int32), + cu_seqlens=infinicore.from_list( + [0, past_length + length], dtype=infinicore.int32 + ), + mamba_init_state_indices=infinicore.from_list( + [initial_row], dtype=infinicore.int32 + ), + mamba_final_state_indices=infinicore.from_list( + [final_row], dtype=infinicore.int32 + ), + sample_all_positions=True, + ) + logits = output["logits"].to(infinicore.device("cpu", 0)) + values = np.empty(logits.shape, dtype=np.float16) + ctypes.memmove(values.ctypes.data, logits.data_ptr(), values.nbytes) + return values + + def test_load_prefill_decode_and_state_isolation(self): + prefill = self._forward([1, 2, 3, 4, 5, 6, 7, 8], 0, 1, 0) + decode = self._forward([9], 1, 1, 8) + parallel = self._forward([10, 11], 0, 2, 0) + + self.assertEqual(prefill.shape, (1, 8, 50288)) + self.assertEqual(decode.shape, (1, 1, 50288)) + self.assertEqual(parallel.shape, (1, 2, 50288)) + self.assertTrue(np.isfinite(prefill).all()) + self.assertTrue(np.isfinite(decode).all()) + self.assertTrue(np.isfinite(parallel).all()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/mamba2/test_scheduler.py b/test/models/mamba2/test_scheduler.py new file mode 100644 index 000000000..8600562db --- /dev/null +++ b/test/models/mamba2/test_scheduler.py @@ -0,0 +1,43 @@ +import unittest + +from infinilm.llm.request import InferenceRequest +from infinilm.llm.sampling_params import SamplingParams +from infinilm.llm.scheduler import Scheduler + + +class Mamba2SchedulerTest(unittest.TestCase): + def test_recurrent_requests_skip_attention_kv_reservation(self): + scheduler = Scheduler( + max_batch_size=4, + num_blocks=128, + block_size=16, + max_num_batched_tokens=4096, + has_mamba_cache=True, + cacheless_state_model=True, + num_mamba_cache_blocks=8, + enable_prefix_caching=False, + ) + requests = [ + InferenceRequest( + request_id=f"mamba2-unit-{index}", + prompt_token_ids=[1] * 512, + sampling_params=SamplingParams(max_tokens=128, ignore_eos=True), + ) + for index in range(4) + ] + for request in requests: + scheduler.add_request(request) + + output = scheduler.schedule() + + self.assertIsNotNone(output) + self.assertEqual(len(output.scheduled_requests), 4) + self.assertTrue(all(not request.block_table for request in requests)) + self.assertTrue(all(not request.slot_mapping for request in requests)) + self.assertEqual( + {request.mamba_cache_index for request in requests}, {1, 2, 3, 4} + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/rwkv5/BENCHMARK.md b/test/models/rwkv5/BENCHMARK.md new file mode 100644 index 000000000..29926dfb8 --- /dev/null +++ b/test/models/rwkv5/BENCHMARK.md @@ -0,0 +1,85 @@ +# RWKV5 NVIDIA Benchmark + +## Environment + +- GPU: NVIDIA GeForce RTX 4090 D, 24 GiB +- Model: RWKV-5-World-0.1B-v1-20230803-ctx4096 +- Precision: BF16 +- Cache mode: paged, 128 blocks, block size 16 +- Sampling: greedy (`top_k=1`), EOS ignored to keep output lengths identical +- Warmup: 3 runs +- Measured runs: 10 per comparison case + +The prompt token IDs are deterministic and identical between the two modes. The +comparison is an architectural ablation: the normal RWKV5 path uses only one +recurrent state row per request, while the forced path exercises the generic +attention-KV scheduler path. RWKV5 itself does not consume attention KV blocks. +This benchmark does not claim to change the generic KV admission or +over-reservation policy used by Transformer models. + +## Metric Definitions + +- TTFT is batch max TTFT: time from submitting the batch until every request has + produced its first token. +- Decode throughput counts only tokens generated after every request has reached + TTFT. Tokens produced by earlier requests while another request is waiting are + not counted twice. +- End-to-end throughput is total requested output tokens divided by total elapsed + generation time. +- P50 is the median. P90 uses the nearest-rank definition. + +## Comparison + +Batch size is 4 and each prompt contains 512 tokens. + +| Output tokens | Mode | TTFT P50 | TTFT P90 | E2E throughput P50 | +|---:|---|---:|---:|---:| +| 32 | Forced attention-KV path | 337.22 ms | 362.65 ms | 254.95 tok/s | +| 32 | Recurrent-only state path | 120.32 ms | 120.98 ms | 362.90 tok/s | +| 128 | Forced attention-KV path | 1017.99 ms | 1065.04 ms | 308.07 tok/s | +| 128 | Recurrent-only state path | 120.08 ms | 120.27 ms | 474.20 tok/s | + +Results: + +- Output 32: TTFT P50 decreased by 64.32% (2.80x), while end-to-end + throughput increased by 42.34%. +- Output 128: TTFT P50 decreased by 88.20% (8.48x), while end-to-end + throughput increased by 53.93%. +- The forced attention-KV path generated 97 or 385 tokens before all requests reached TTFT. This + means the first three requests completed before the fourth request started. + The recurrent-only path generated exactly four tokens at the same boundary, one + per request, confirming that RWKV5 does not request attention KV blocks. +- Peak GPU memory was 1203 MiB in both modes. The architectural path avoids + attention-KV allocation for RWKV5; it does not change the generic KV + reservation policy or the RWKV state pool allocation. + +## Reproduction + +Recurrent-only state path: + +```bash +RWKV5_MODEL_PATH=/models/RWKV-5-World-0.1B-InfiniLM \ +python test/models/rwkv5/benchmark.py \ + --cache-type paged --num-blocks 128 \ + --batch-sizes 4 --input-lens 512 --output-lens 32,128 \ + --warmup 3 --runs 10 +``` + +Forced attention-KV ablation: + +```bash +RWKV5_MODEL_PATH=/models/RWKV-5-World-0.1B-InfiniLM \ +python test/models/rwkv5/benchmark.py \ + --cache-type paged --num-blocks 128 \ + --batch-sizes 4 --input-lens 512 --output-lens 32,128 \ + --warmup 3 --runs 10 --force-attention-kv +``` + +Raw results: + +- `results/rwkv5-p50-optimized-v2.jsonl` +- `results/rwkv5-p50-legacy-v2.jsonl` +- `results/rwkv5-full-optimized-p50.jsonl` + +The full optimized matrix covers batch sizes 1, 2, and 4; input lengths 32, +128, and 512; and output lengths 32 and 128, with five measured runs per case. diff --git a/test/models/rwkv5/benchmark.py b/test/models/rwkv5/benchmark.py new file mode 100644 index 000000000..bc01c12f4 --- /dev/null +++ b/test/models/rwkv5/benchmark.py @@ -0,0 +1,205 @@ +"""Reproducible single-GPU RWKV5 scheduler benchmark.""" + +import argparse +import json +import math +import os +import statistics +import time + +import torch + +from infinilm.llm.llm import LLM +from infinilm.llm.request import InferenceRequest +from infinilm.llm.sampling_params import SamplingParams + + +def parse_int_list(value: str) -> list[int]: + return [int(item) for item in value.split(",") if item.strip()] + + +def make_prompt_tokens(tokenizer, length: int) -> list[int]: + seed = tokenizer.encode( + "The quick brown fox jumps over the lazy dog. " + "InfiniLM RWKV5 scheduler benchmark. " + ) + return (seed * ((length + len(seed) - 1) // len(seed)))[:length] + + +def run_case(engine, tokenizer, batch_size: int, input_len: int, output_len: int): + prompt_tokens = make_prompt_tokens(tokenizer, input_len) + sampling = SamplingParams( + max_tokens=output_len, + temperature=1.0, + top_p=1.0, + top_k=1, + ignore_eos=True, + ) + requests = [] + for index in range(batch_size): + request = InferenceRequest( + request_id=f"rwkv5-bench-{batch_size}-{input_len}-{index}", + prompt_token_ids=prompt_tokens, + sampling_params=sampling, + eos_token_ids=engine.eos_token_ids, + ) + requests.append(request) + engine.add_request(request) + + torch.cuda.synchronize() + start = time.perf_counter() + while any(request.get_num_generated_tokens() == 0 for request in requests): + did_work, _ = engine.step() + if not did_work: + raise RuntimeError("scheduler made no progress during prefill") + torch.cuda.synchronize() + ttft = time.perf_counter() - start + generated_at_ttft = [ + request.get_num_generated_tokens() for request in requests + ] + + decode_start = time.perf_counter() + while any(not request.is_finished() for request in requests): + did_work, _ = engine.step() + if not did_work: + raise RuntimeError("scheduler made no progress during decode") + torch.cuda.synchronize() + decode_time = time.perf_counter() - decode_start + + generated = [request.get_num_generated_tokens() for request in requests] + expected = [output_len] * batch_size + if generated != expected: + raise AssertionError(f"unexpected output lengths: {generated} != {expected}") + + remaining_tokens = [output_len - count for count in generated_at_ttft] + decode_tokens = sum(remaining_tokens) + decode_steps = max(remaining_tokens, default=0) + return { + "batch_size": batch_size, + "input_len": input_len, + "output_len": output_len, + "ttft_ms": ttft * 1000.0, + "tokens_generated_before_all_ttft": sum(generated_at_ttft), + "decode_time_ms": decode_time * 1000.0, + "decode_itl_ms": decode_time * 1000.0 / max(decode_steps, 1), + "decode_throughput_tok_s": ( + decode_tokens / decode_time if decode_tokens else 0.0 + ), + "end_to_end_throughput_tok_s": batch_size * output_len / (ttft + decode_time), + } + + +def percentile(values: list[float], percent: float) -> float: + ordered = sorted(values) + index = max(0, math.ceil(percent * len(ordered)) - 1) + return ordered[index] + + +def summarize(samples: list[dict]) -> dict: + summary = { + key: samples[0][key] + for key in ("batch_size", "input_len", "output_len") + } + summary["runs"] = len(samples) + for metric in ( + "ttft_ms", + "decode_itl_ms", + "decode_throughput_tok_s", + "end_to_end_throughput_tok_s", + ): + values = [sample[metric] for sample in samples] + summary[f"{metric}_p50"] = statistics.median(values) + summary[f"{metric}_p90"] = percentile(values, 0.9) + return summary + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", default=os.environ.get("RWKV5_MODEL_PATH"), required=False + ) + parser.add_argument("--batch-sizes", default="1,2,4", type=parse_int_list) + parser.add_argument("--input-lens", default="32,128,512", type=parse_int_list) + parser.add_argument("--output-lens", default="32,128", type=parse_int_list) + parser.add_argument("--cache-type", choices=("paged", "static"), default="paged") + parser.add_argument("--num-blocks", type=int, default=128) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--runs", type=int, default=5) + parser.add_argument( + "--force-attention-kv", + "--legacy-kv-reservation", + dest="force_attention_kv", + action="store_true", + help="Ablate the pure-recurrent path by forcing attention KV scheduling", + ) + args = parser.parse_args() + if not args.model: + parser.error("--model or RWKV5_MODEL_PATH is required") + if args.cache_type == "static" and args.batch_sizes != [1]: + parser.error( + "static cache uses the single-request scheduler; use --batch-sizes 1" + ) + if args.runs < 1: + parser.error("--runs must be at least 1") + if args.force_attention_kv and args.cache_type != "paged": + parser.error("--force-attention-kv requires --cache-type paged") + + model = LLM( + model_path=args.model, + device="cuda", + dtype="bfloat16", + tensor_parallel_size=1, + cache_type=args.cache_type, + max_batch_size=max(args.batch_sizes), + max_tokens=max(args.output_lens), + num_blocks=args.num_blocks, + block_size=16, + enable_graph=False, + attn_backend="paged-attn", + weight_load_mode="sync", + enable_prefix_caching=False, + ) + try: + engine = model.engine + tokenizer = engine.tokenizer + if args.force_attention_kv: + engine.scheduler.cacheless_state_model = False + for _ in range(args.warmup): + run_case(engine, tokenizer, 1, min(args.input_lens), min(args.output_lens)) + + print( + json.dumps( + { + "type": "metadata", + "cache_type": args.cache_type, + "num_blocks": args.num_blocks, + "force_attention_kv": args.force_attention_kv, + "runs": args.runs, + } + ) + ) + for batch_size in args.batch_sizes: + for input_len in args.input_lens: + for output_len in args.output_lens: + samples = [] + for run_index in range(args.runs): + result = run_case( + engine, tokenizer, batch_size, input_len, output_len + ) + samples.append(result) + print( + json.dumps( + {"type": "sample", "run": run_index + 1, **result} + ), + flush=True, + ) + print( + json.dumps({"type": "summary", **summarize(samples)}), + flush=True, + ) + finally: + model.close() + + +if __name__ == "__main__": + main() diff --git a/test/models/rwkv5/test_adaptation.py b/test/models/rwkv5/test_adaptation.py new file mode 100644 index 000000000..162d4d251 --- /dev/null +++ b/test/models/rwkv5/test_adaptation.py @@ -0,0 +1,135 @@ +import json +import tempfile +import unittest +from pathlib import Path + +import torch + +from infinilm.infer_engine import InferEngine, model_uses_mamba_cache +from infinilm.modeling_utils import _remap_rwkv5 +from infinilm.processors import AutoInfinilmProcessor +from infinilm.processors.rwkv5_processor import RWKV5Processor, RWKVWorldTokenizer +from scripts.convert_rwkv5_checkpoint import _infer_config + + +class RWKV5CheckpointTest(unittest.TestCase): + def _state_dict(self, time_decay, *, gated=False): + state_dict = { + "emb.weight": torch.zeros(16, 8), + "blocks.0.att.time_decay": time_decay, + "blocks.0.att.time_first": torch.log( + torch.linspace(0.5, 0.75, time_decay.shape[0]) + ), + "blocks.0.ffn.key.weight": torch.zeros(16, 8), + } + if gated: + state_dict["blocks.0.att.gate.weight"] = torch.zeros(8, 8) + state_dict["blocks.0.att.time_mix_g"] = torch.zeros(1, 1, 8) + return state_dict + + def test_infers_rwkv50_scalar_per_head_layout(self): + config = _infer_config( + self._state_dict(torch.tensor([-2.0, -3.0])), context_length=1024 + ) + + self.assertEqual(config["rwkv_version"], "5.0") + self.assertEqual(config["num_attention_heads"], 2) + self.assertEqual(config["head_dim"], 4) + self.assertFalse(config["use_attention_gate"]) + + def test_infers_rwkv52_per_channel_layout(self): + config = _infer_config( + self._state_dict(torch.zeros(2, 4), gated=True), context_length=2048 + ) + + self.assertEqual(config["rwkv_version"], "5.2") + self.assertEqual(config["num_attention_heads"], 2) + self.assertEqual(config["head_dim"], 4) + self.assertTrue(config["use_attention_gate"]) + + def test_infers_rwkv51_gated_scalar_per_head_layout(self): + config = _infer_config( + self._state_dict(torch.zeros(2), gated=True), context_length=2048 + ) + + self.assertEqual(config["rwkv_version"], "5.1") + self.assertTrue(config["use_attention_gate"]) + + def test_remaps_and_expands_rwkv50_time_parameters(self): + state_dict = self._state_dict(torch.tensor([-2.0, -3.0])) + remapped = _remap_rwkv5( + state_dict, + {"num_attention_heads": 2, "head_dim": 4}, + ) + + self.assertEqual(remapped["model.blocks.0.att.time_decay"].shape, (2, 4)) + torch.testing.assert_close( + remapped["model.blocks.0.att.time_decay"][:, 0], + torch.tensor([-2.0, -3.0]), + ) + expected_first = torch.linspace(0.5, 0.75, 2) + torch.testing.assert_close( + remapped["model.blocks.0.att.time_faaaa"][:, 0], expected_first + ) + self.assertEqual(remapped["model.embeddings.weight"].shape, (16, 8)) + + def test_rejects_invalid_time_parameter_size(self): + with self.assertRaisesRegex(ValueError, "expected 2 or 8"): + _remap_rwkv5( + {"blocks.0.att.time_decay": torch.zeros(3)}, + {"num_attention_heads": 2, "head_dim": 4}, + ) + + +class RWKV5ProcessorTest(unittest.TestCase): + VOCAB = """1 b'a' 1 +2 b'b' 1 +3 b'ab' 2 +4 '\u4f60' 3 +""" + + def test_world_tokenizer_prefers_longest_byte_match(self): + with tempfile.TemporaryDirectory() as directory: + vocab_path = Path(directory) / "rwkv_vocab_v20230424.txt" + vocab_path.write_text(self.VOCAB, encoding="utf-8") + tokenizer = RWKVWorldTokenizer(vocab_path) + + self.assertEqual(tokenizer.encode("ab\u4f60"), [3, 4]) + self.assertEqual(tokenizer.decode([3, 4]), "ab\u4f60") + + def test_auto_processor_uses_registered_rwkv5_processor(self): + with tempfile.TemporaryDirectory() as directory: + model_dir = Path(directory) + (model_dir / "config.json").write_text( + json.dumps({"model_type": "rwkv5"}), encoding="utf-8" + ) + (model_dir / "rwkv_vocab_v20230424.txt").write_text( + self.VOCAB, encoding="utf-8" + ) + + processor = AutoInfinilmProcessor.from_pretrained(str(model_dir)) + + self.assertIsInstance(processor, RWKV5Processor) + self.assertEqual(processor("ab")["input_ids"], [3]) + + def test_rwkv5_uses_request_state_cache(self): + self.assertTrue(model_uses_mamba_cache({"model_type": "rwkv5"})) + + def test_zero_is_a_valid_eos_token_id(self): + engine = type("EngineConfigStub", (), {})() + engine.hf_generation_config = {"eos_token_id": 0} + engine.hf_config = {"eos_token_id": 7} + + self.assertEqual(InferEngine.eos_token_id.fget(engine), [0]) + + def test_eos_is_not_decoded_as_a_byte_token(self): + with tempfile.TemporaryDirectory() as directory: + vocab_path = Path(directory) / "rwkv_vocab_v20230424.txt" + vocab_path.write_text(self.VOCAB, encoding="utf-8") + tokenizer = RWKVWorldTokenizer(vocab_path) + + self.assertEqual(tokenizer.decode([tokenizer.eos_token_id]), "") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/rwkv5/test_correctness.py b/test/models/rwkv5/test_correctness.py new file mode 100644 index 000000000..6e58ca44e --- /dev/null +++ b/test/models/rwkv5/test_correctness.py @@ -0,0 +1,345 @@ +import json +import os +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from safetensors.torch import save_file + +import infinicore +from infinilm.cache import PagedKVCacheConfig +from infinilm.distributed import DistConfig +from infinilm.infer_engine import InferEngine +from infinilm.modeling_utils import load_model_state_dict_by_file + + +RUN_GPU_TESTS = os.getenv("INFINILM_RUN_GPU_TESTS") == "1" + + +def _make_tiny_checkpoint(seed=20260920): + generator = torch.Generator().manual_seed(seed) + vocab_size = 32 + hidden_size = 8 + intermediate_size = 16 + num_heads = 2 + head_dim = hidden_size // num_heads + num_layers = 2 + + def randn(*shape, scale=0.08): + return torch.randn(*shape, generator=generator) * scale + + state = { + "emb.weight": randn(vocab_size, hidden_size), + "blocks.0.ln0.weight": 1 + randn(hidden_size, scale=0.02), + "blocks.0.ln0.bias": randn(hidden_size, scale=0.02), + "ln_out.weight": 1 + randn(hidden_size, scale=0.02), + "ln_out.bias": randn(hidden_size, scale=0.02), + "head.weight": randn(vocab_size, hidden_size), + } + for layer_idx in range(num_layers): + block = f"blocks.{layer_idx}" + state.update( + { + f"{block}.ln1.weight": 1 + randn(hidden_size, scale=0.02), + f"{block}.ln1.bias": randn(hidden_size, scale=0.02), + f"{block}.ln2.weight": 1 + randn(hidden_size, scale=0.02), + f"{block}.ln2.bias": randn(hidden_size, scale=0.02), + f"{block}.att.time_mix_k": torch.rand( + 1, 1, hidden_size, generator=generator + ), + f"{block}.att.time_mix_v": torch.rand( + 1, 1, hidden_size, generator=generator + ), + f"{block}.att.time_mix_r": torch.rand( + 1, 1, hidden_size, generator=generator + ), + f"{block}.att.time_mix_g": torch.rand( + 1, 1, hidden_size, generator=generator + ), + f"{block}.att.time_decay": randn( + num_heads, head_dim, scale=0.4 + ) + - 2.0, + f"{block}.att.time_faaaa": torch.rand( + num_heads, head_dim, generator=generator + ), + f"{block}.att.key.weight": randn(hidden_size, hidden_size), + f"{block}.att.value.weight": randn(hidden_size, hidden_size), + f"{block}.att.receptance.weight": randn( + hidden_size, hidden_size + ), + f"{block}.att.gate.weight": randn(hidden_size, hidden_size), + f"{block}.att.output.weight": randn(hidden_size, hidden_size), + f"{block}.att.ln_x.weight": 1 + + randn(hidden_size, scale=0.02), + f"{block}.att.ln_x.bias": randn(hidden_size, scale=0.02), + f"{block}.ffn.time_mix_k": torch.rand( + 1, 1, hidden_size, generator=generator + ), + f"{block}.ffn.time_mix_r": torch.rand( + 1, 1, hidden_size, generator=generator + ), + f"{block}.ffn.key.weight": randn( + intermediate_size, hidden_size + ), + f"{block}.ffn.value.weight": randn( + hidden_size, intermediate_size + ), + f"{block}.ffn.receptance.weight": randn( + hidden_size, hidden_size + ), + } + ) + + config = { + "architectures": ["RWKV5ForCausalLM"], + "model_type": "rwkv5", + "rwkv_version": "5.2", + "vocab_size": vocab_size, + "hidden_size": hidden_size, + "intermediate_size": intermediate_size, + "num_hidden_layers": num_layers, + "num_attention_heads": num_heads, + "head_dim": head_dim, + "max_position_embeddings": 128, + "layer_norm_eps": 1e-5, + "group_norm_eps": 64e-5, + "use_attention_gate": True, + "torch_dtype": "float32", + "bos_token_id": 0, + "eos_token_id": 0, + "pad_token_id": 0, + } + return state, config + + +class TorchRWKV5Reference: + def __init__(self, weights, config): + self.weights = weights + self.config = config + self.states = {} + + def _new_state(self): + h = self.config["num_attention_heads"] + n = self.config["head_dim"] + c = self.config["hidden_size"] + return [ + { + "att_prev": torch.zeros(c), + "wkv": torch.zeros(h, n, n), + "ffn_prev": torch.zeros(c), + } + for _ in range(self.config["num_hidden_layers"]) + ] + + @staticmethod + def _mix(previous, current, amount): + return torch.lerp(previous, current, amount) + + def forward(self, token_ids, state_id): + state = self.states.setdefault(state_id, self._new_state()) + w = self.weights + c = self.config["hidden_size"] + h = self.config["num_attention_heads"] + n = self.config["head_dim"] + eps = self.config["layer_norm_eps"] + group_eps = self.config["group_norm_eps"] + + x = F.embedding(torch.tensor(token_ids), w["emb.weight"]) + x = F.layer_norm( + x, + (c,), + w["blocks.0.ln0.weight"], + w["blocks.0.ln0.bias"], + eps, + ) + + for layer_idx, layer_state in enumerate(state): + block = f"blocks.{layer_idx}" + att = f"{block}.att" + ffn = f"{block}.ffn" + + xx = F.layer_norm( + x, (c,), w[f"{block}.ln1.weight"], w[f"{block}.ln1.bias"], eps + ) + previous = torch.cat([layer_state["att_prev"][None], xx[:-1]], dim=0) + layer_state["att_prev"] = xx[-1].clone() + k = F.linear( + self._mix(previous, xx, w[f"{att}.time_mix_k"].reshape(c)), + w[f"{att}.key.weight"], + ) + v = F.linear( + self._mix(previous, xx, w[f"{att}.time_mix_v"].reshape(c)), + w[f"{att}.value.weight"], + ) + r = F.linear( + self._mix(previous, xx, w[f"{att}.time_mix_r"].reshape(c)), + w[f"{att}.receptance.weight"], + ) + g = F.silu( + F.linear( + self._mix(previous, xx, w[f"{att}.time_mix_g"].reshape(c)), + w[f"{att}.gate.weight"], + ) + ) + + decay = torch.exp(-torch.exp(w[f"{att}.time_decay"]))[:, :, None] + first = w[f"{att}.time_faaaa"][:, :, None] + time_out = [] + for token_idx in range(len(token_ids)): + rt = r[token_idx].reshape(h, n) + kt = k[token_idx].reshape(h, n) + vt = v[token_idx].reshape(h, n) + kv = kt[:, :, None] * vt[:, None, :] + yt = torch.matmul( + rt[:, None, :], first * kv + layer_state["wkv"] + ).squeeze(1) + time_out.append(yt.reshape(c)) + layer_state["wkv"] = kv + decay * layer_state["wkv"] + time_out = torch.stack(time_out) + time_out = F.group_norm( + time_out, + num_groups=h, + weight=w[f"{att}.ln_x.weight"], + bias=w[f"{att}.ln_x.bias"], + eps=group_eps, + ) + x = x + F.linear(time_out * g, w[f"{att}.output.weight"]) + + xx = F.layer_norm( + x, (c,), w[f"{block}.ln2.weight"], w[f"{block}.ln2.bias"], eps + ) + previous = torch.cat([layer_state["ffn_prev"][None], xx[:-1]], dim=0) + layer_state["ffn_prev"] = xx[-1].clone() + k = F.linear( + self._mix(previous, xx, w[f"{ffn}.time_mix_k"].reshape(c)), + w[f"{ffn}.key.weight"], + ) + r = F.linear( + self._mix(previous, xx, w[f"{ffn}.time_mix_r"].reshape(c)), + w[f"{ffn}.receptance.weight"], + ) + x = x + torch.sigmoid(r) * F.linear( + torch.relu(k).square(), w[f"{ffn}.value.weight"] + ) + + x = F.layer_norm( + x, (c,), w["ln_out.weight"], w["ln_out.bias"], eps + ) + return F.linear(x, w["head.weight"]) + + +@unittest.skipUnless( + RUN_GPU_TESTS and torch.cuda.is_available(), + "set INFINILM_RUN_GPU_TESTS=1 on an NVIDIA host", +) +class RWKV5GPUCorrectnessTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.temp_dir = tempfile.TemporaryDirectory() + model_dir = Path(cls.temp_dir.name) + cls.weights, cls.config = _make_tiny_checkpoint() + save_file(cls.weights, model_dir / "model.safetensors", metadata={"format": "pt"}) + (model_dir / "config.json").write_text( + json.dumps(cls.config), encoding="utf-8" + ) + (model_dir / "generation_config.json").write_text( + json.dumps({"eos_token_id": 0}), encoding="utf-8" + ) + + cls.engine = InferEngine( + str(model_dir), + device=infinicore.device("cuda", 0), + distributed_config=DistConfig(1), + cache_config=PagedKVCacheConfig( + num_blocks=16, block_size=16, max_batch_size=4 + ), + attention_backend="paged-attn", + weight_load_mode="sync", + ) + load_model_state_dict_by_file( + cls.engine, str(model_dir), dtype=infinicore.float32 + ) + + @classmethod + def tearDownClass(cls): + del cls.engine + cls.temp_dir.cleanup() + + def _forward(self, requests, initial_rows, final_rows, past_lengths): + flat_tokens = [token for request in requests for token in request] + offsets = [0] + for request in requests: + offsets.append(offsets[-1] + len(request)) + total_lengths = [past + len(request) for past, request in zip(past_lengths, requests)] + cu_seqlens = [0] + for total_length in total_lengths: + cu_seqlens.append(cu_seqlens[-1] + total_length) + + output = self.engine.forward_raw( + infinicore.from_list([flat_tokens], dtype=infinicore.int64), + position_ids=infinicore.from_list( + [ + position + for past, request in zip(past_lengths, requests) + for position in range(past, past + len(request)) + ], + dtype=infinicore.int64, + ), + past_kv_lengths=infinicore.from_list( + past_lengths, dtype=infinicore.int32 + ), + total_kv_lengths=infinicore.from_list( + total_lengths, dtype=infinicore.int32 + ), + input_offsets=infinicore.from_list(offsets, dtype=infinicore.int32), + cu_seqlens=infinicore.from_list(cu_seqlens, dtype=infinicore.int32), + mamba_init_state_indices=infinicore.from_list( + initial_rows, dtype=infinicore.int32 + ), + mamba_final_state_indices=infinicore.from_list( + final_rows, dtype=infinicore.int32 + ), + sample_all_positions=True, + ) + return torch.from_numpy(np.asarray(output["logits"].to_numpy())).reshape( + len(flat_tokens), -1 + ) + + def test_prefill_decode_and_concurrent_state_isolation(self): + reference = TorchRWKV5Reference(self.weights, self.config) + + prompt = [3, 7, 11, 5] + actual = self._forward([prompt], [0], [1], [0]) + expected = reference.forward(prompt, state_id=1) + torch.testing.assert_close(actual, expected, atol=3e-4, rtol=3e-4) + + actual = self._forward([[9]], [1], [1], [len(prompt)]) + expected = reference.forward([9], state_id=1) + torch.testing.assert_close(actual, expected, atol=3e-4, rtol=3e-4) + + prompts = [[2, 4, 6], [13, 17]] + actual = self._forward(prompts, [0, 0], [2, 3], [0, 0]) + expected = torch.cat( + [ + reference.forward(prompts[0], state_id=2), + reference.forward(prompts[1], state_id=3), + ] + ) + torch.testing.assert_close(actual, expected, atol=3e-4, rtol=3e-4) + + actual = self._forward([[8], [19]], [2, 3], [2, 3], [3, 2]) + expected = torch.cat( + [ + reference.forward([8], state_id=2), + reference.forward([19], state_id=3), + ] + ) + torch.testing.assert_close(actual, expected, atol=3e-4, rtol=3e-4) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/rwkv5/test_real_model.py b/test/models/rwkv5/test_real_model.py new file mode 100644 index 000000000..fcda215cf --- /dev/null +++ b/test/models/rwkv5/test_real_model.py @@ -0,0 +1,82 @@ +import os +import unittest + +import torch + +from infinilm.llm.llm import LLM +from infinilm.llm.sampling_params import SamplingParams + + +MODEL_PATH = os.getenv("RWKV5_MODEL_PATH") +RUN_GPU_TESTS = os.getenv("INFINILM_RUN_GPU_TESTS") == "1" + + +@unittest.skipUnless( + RUN_GPU_TESTS and MODEL_PATH and torch.cuda.is_available(), + "set INFINILM_RUN_GPU_TESTS=1 and RWKV5_MODEL_PATH on an NVIDIA host", +) +class RWKV5RealModelTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.model = LLM( + model_path=MODEL_PATH, + device="cuda", + dtype="bfloat16", + tensor_parallel_size=1, + cache_type="paged", + max_batch_size=4, + max_tokens=24, + num_blocks=64, + block_size=16, + temperature=1.0, + top_p=1.0, + top_k=1, + enable_graph=False, + attn_backend="paged-attn", + weight_load_mode="sync", + enable_prefix_caching=False, + ) + cls.sampling = SamplingParams( + max_tokens=24, + temperature=1.0, + top_p=1.0, + top_k=1, + ) + + @classmethod + def tearDownClass(cls): + cls.model.close() + + @staticmethod + def _conversations(): + return [ + [{"role": "user", "content": "What is the capital of France?"}], + [{"role": "user", "content": "Write the first three prime numbers."}], + ] + + def test_known_answer_and_batched_state_isolation(self): + conversations = self._conversations() + batched = self.model.chat( + messages=conversations, + sampling_params=self.sampling, + use_tqdm=False, + ) + sequential = [ + self.model.chat( + messages=conversation, + sampling_params=self.sampling, + use_tqdm=False, + )[0] + for conversation in conversations + ] + + self.assertIn("Paris", batched[0].outputs[0].text) + for batched_output, sequential_output in zip(batched, sequential): + self.assertEqual( + batched_output.outputs[0].token_ids, + sequential_output.outputs[0].token_ids, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/models/rwkv5/test_scheduler.py b/test/models/rwkv5/test_scheduler.py new file mode 100644 index 000000000..5144870ba --- /dev/null +++ b/test/models/rwkv5/test_scheduler.py @@ -0,0 +1,87 @@ +import unittest + +from infinilm.llm.request import InferenceRequest +from infinilm.llm.sampling_params import SamplingParams +from infinilm.llm.scheduler import Scheduler + + +class RWKV5SchedulerTest(unittest.TestCase): + def test_cacheless_state_model_does_not_reserve_attention_blocks(self): + scheduler = Scheduler( + max_batch_size=4, + num_blocks=128, + block_size=16, + max_num_batched_tokens=4096, + has_mamba_cache=True, + cacheless_state_model=True, + num_mamba_cache_blocks=8, + enable_prefix_caching=False, + ) + requests = [ + InferenceRequest( + request_id=f"rwkv5-unit-{index}", + prompt_token_ids=[1] * 512, + sampling_params=SamplingParams(max_tokens=128, ignore_eos=True), + ) + for index in range(4) + ] + for request in requests: + scheduler.add_request(request) + + output = scheduler.schedule() + + self.assertIsNotNone(output) + self.assertEqual(len(output.scheduled_requests), 4) + self.assertTrue(all(not request.block_table for request in requests)) + self.assertTrue(all(not request.slot_mapping for request in requests)) + self.assertTrue(all(request.mamba_cache_index is not None for request in requests)) + + def test_default_scheduler_still_reserves_attention_blocks(self): + scheduler = Scheduler( + max_batch_size=1, + num_blocks=128, + block_size=16, + max_num_batched_tokens=1024, + enable_prefix_caching=False, + ) + request = InferenceRequest( + request_id="transformer-unit", + prompt_token_ids=[1] * 32, + sampling_params=SamplingParams(max_tokens=32, ignore_eos=True), + ) + scheduler.add_request(request) + + output = scheduler.schedule() + + self.assertIsNotNone(output) + self.assertEqual(output.scheduled_requests, [request]) + self.assertTrue(request.block_table) + self.assertEqual(len(request.slot_mapping), 32) + + def test_mamba_state_cache_does_not_change_attention_kv_path(self): + scheduler = Scheduler( + max_batch_size=1, + num_blocks=128, + block_size=16, + max_num_batched_tokens=1024, + has_mamba_cache=True, + cacheless_state_model=False, + num_mamba_cache_blocks=8, + enable_prefix_caching=False, + ) + request = InferenceRequest( + request_id="hybrid-unit", + prompt_token_ids=[1] * 32, + sampling_params=SamplingParams(max_tokens=32, ignore_eos=True), + ) + scheduler.add_request(request) + + output = scheduler.schedule() + + self.assertIsNotNone(output) + self.assertTrue(request.block_table) + self.assertEqual(request.mamba_cache_index, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/xmake.lua b/xmake.lua index dd04bbbee..cc029c34d 100644 --- a/xmake.lua +++ b/xmake.lua @@ -1,5 +1,3 @@ -add_requires("pybind11") - set_toolchains("gcc") option("cxx11-abi") @@ -20,13 +18,16 @@ end -- Add spdlog from third_party directory add_includedirs("third_party/spdlog/include") add_includedirs("third_party/json/single_include/") +add_includedirs("/usr/include/python3.12") target("_infinilm") - add_packages("pybind11") set_default(false) - add_rules("python.module", {soabi = true}) + -- The python.module rule is not shipped by the minimal xmake package used + -- on the NVIDIA benchmark host. A shared library is importable by Python + -- on Linux and keeps the build self-contained. set_languages("cxx17") set_kind("shared") + set_prefixname("") local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") @@ -36,7 +37,8 @@ target("_infinilm") -- spdlog is already included globally via add_includedirs at the top add_linkdirs(INFINI_ROOT.."/lib") - add_links("infinicore_cpp_api", "infiniop", "infinirt", "infiniccl") + add_links("infinicore_cpp_api", "infiniop", "infinirt", "infiniccl", "fmt") + add_shflags("-Wl,--no-as-needed", "-lfmt", "-Wl,--as-needed", { force = true }) -- Add C++ sources add_files("csrc/**.cpp")