From 63b5418e579df8b95f65a4b166b7e5f090c45b29 Mon Sep 17 00:00:00 2001 From: liuzhenlong Date: Sat, 4 Jul 2026 17:48:08 +0800 Subject: [PATCH] feat: add RWKV-7-World model support for CUDA inference --- tests/core/framework/hf_model_loader_test.cpp | 41 ++ tests/core/framework/tokenizer/CMakeLists.txt | 11 + .../tokenizer/rwkv_tokenizer_test.cpp | 165 +++++ tests/core/layers/cuda/CMakeLists.txt | 21 + .../layers/cuda/rwkv7_decoder_layer_test.cpp | 250 +++++++ tools/convert_rwkv7_world.py | 319 +++++++++ xllm/core/framework/tokenizer/CMakeLists.txt | 2 + .../framework/tokenizer/rwkv_tokenizer.cpp | 604 +++++++++++++++++ .../core/framework/tokenizer/rwkv_tokenizer.h | 66 ++ .../core/framework/tokenizer/tokenizer_args.h | 4 +- .../framework/tokenizer/tokenizer_factory.cpp | 6 +- .../framework/tokenizer/tokenizer_factory.h | 1 + xllm/core/layers/CMakeLists.txt | 2 + xllm/core/layers/rwkv7_decoder_layer.cpp | 614 ++++++++++++++++++ xllm/core/layers/rwkv7_decoder_layer.h | 246 +++++++ xllm/models/llm/rwkv7.h | 349 ++++++++++ xllm/models/models.h | 1 + 17 files changed, 2699 insertions(+), 3 deletions(-) create mode 100644 tests/core/framework/tokenizer/rwkv_tokenizer_test.cpp create mode 100644 tests/core/layers/cuda/rwkv7_decoder_layer_test.cpp create mode 100755 tools/convert_rwkv7_world.py create mode 100644 xllm/core/framework/tokenizer/rwkv_tokenizer.cpp create mode 100644 xllm/core/framework/tokenizer/rwkv_tokenizer.h create mode 100644 xllm/core/layers/rwkv7_decoder_layer.cpp create mode 100644 xllm/core/layers/rwkv7_decoder_layer.h create mode 100644 xllm/models/llm/rwkv7.h diff --git a/tests/core/framework/hf_model_loader_test.cpp b/tests/core/framework/hf_model_loader_test.cpp index 5d13e7e025..7f5bc1d3f3 100644 --- a/tests/core/framework/hf_model_loader_test.cpp +++ b/tests/core/framework/hf_model_loader_test.cpp @@ -343,6 +343,47 @@ TEST(HFModelLoaderTest, MiMoMtpModelArgsDefaults) { EXPECT_EQ(args.eos_token_id(), 151643); EXPECT_EQ(args.stop_token_ids(), std::unordered_set({151643})); } + +TEST(HFModelLoaderTest, Rwkv7ModelArgsFromConvertedConfig) { + auto loader = ModelRegistry::get_model_args_loader("rwkv7"); + ASSERT_NE(loader, nullptr); + + JsonReader reader; + ASSERT_TRUE(reader.parse_text(R"json( + { + "model_type": "rwkv7", + "torch_dtype": "float16", + "vocab_size": 65536, + "hidden_size": 768, + "num_hidden_layers": 12, + "head_size": 64, + "intermediate_size": 3072, + "num_attention_heads": 12, + "layer_norm_eps": 1e-5, + "max_position_embeddings": 4096, + "bos_token_id": 0, + "eos_token_id": 0 + } + )json")); + + ModelArgs args; + ASSERT_TRUE(loader(reader, &args)); + EXPECT_EQ(args.model_type(), "rwkv7"); + EXPECT_EQ(args.vocab_size(), 65536); + EXPECT_EQ(args.hidden_size(), 768); + EXPECT_EQ(args.n_layers(), 12); + EXPECT_EQ(args.head_dim(), 64); + EXPECT_EQ(args.intermediate_size(), 3072); + EXPECT_EQ(args.n_heads(), 12); + EXPECT_EQ(args.max_position_embeddings(), 4096); + EXPECT_FLOAT_EQ(args.layer_norm_eps(), 1e-5f); + EXPECT_EQ(args.linear_conv_kernel_dim(), 2); + EXPECT_EQ(args.full_attention_interval(), 13); + EXPECT_TRUE(has_linear_attention_layers(args)); + ASSERT_EQ(args.layer_types().size(), 12U); + EXPECT_EQ(args.layer_types().front(), "rwkv7"); + EXPECT_EQ(args.stop_token_ids(), std::unordered_set({0})); +} #endif // USE_CUDA #if defined(USE_DCU) diff --git a/tests/core/framework/tokenizer/CMakeLists.txt b/tests/core/framework/tokenizer/CMakeLists.txt index a12dc4fe58..672582690a 100644 --- a/tests/core/framework/tokenizer/CMakeLists.txt +++ b/tests/core/framework/tokenizer/CMakeLists.txt @@ -10,3 +10,14 @@ cc_test( glog::glog GTest::gtest_main ) + +cc_test( + NAME + rwkv_tokenizer_test + SRCS + rwkv_tokenizer_test.cpp + DEPS + :tokenizer + glog::glog + GTest::gtest_main +) diff --git a/tests/core/framework/tokenizer/rwkv_tokenizer_test.cpp b/tests/core/framework/tokenizer/rwkv_tokenizer_test.cpp new file mode 100644 index 0000000000..244a8eaa2b --- /dev/null +++ b/tests/core/framework/tokenizer/rwkv_tokenizer_test.cpp @@ -0,0 +1,165 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "framework/tokenizer/rwkv_tokenizer.h" + +#include + +#include +#include +#include +#include + +#include "framework/tokenizer/tokenizer_args.h" + +namespace xllm { +namespace { + +constexpr char kRwkvVocabFile[] = "rwkv_vocab_test.txt"; + +bool WriteTestRwkvVocab(const std::filesystem::path& filepath) { + std::ofstream file(filepath); + if (!file.is_open()) { + return false; + } + // Minimal trie vocab in the official rwkv_vocab_v20230424.txt line format: + // + file << "0 '<|bos|>' 7\n"; + file << "1 '<|eos|>' 7\n"; + file << "2 'hello' 5\n"; + file << "3 ' world' 6\n"; + file << "4 'test' 4\n"; + return file.good(); +} + +TokenizerArgs MakeRwkvArgs() { + TokenizerArgs args; + args.tokenizer_type() = "rwkv"; + args.vocab_file() = kRwkvVocabFile; + return args; +} + +} // namespace + +class RwkvTokenizerTest : public ::testing::Test { + protected: + void SetUp() override { + test_dir_ = std::filesystem::temp_directory_path() / "rwkv_tokenizer_test"; + std::filesystem::create_directories(test_dir_); + + vocab_path_ = test_dir_ / kRwkvVocabFile; + ASSERT_TRUE(WriteTestRwkvVocab(vocab_path_)); + } + + void TearDown() override { + if (std::filesystem::exists(test_dir_)) { + std::filesystem::remove_all(test_dir_); + } + } + + std::filesystem::path test_dir_; + std::filesystem::path vocab_path_; +}; + +TEST_F(RwkvTokenizerTest, EncodeDecodeRoundTrip) { + RwkvTokenizer tokenizer(test_dir_.string(), MakeRwkvArgs()); + EXPECT_EQ(tokenizer.vocab_size(), 5U); + + std::vector ids; + ASSERT_TRUE( + tokenizer.encode("hello world", &ids, /*add_special_tokens=*/false)); + EXPECT_EQ(ids, std::vector({2, 3})); + + const std::string decoded = + tokenizer.decode(ids, /*skip_special_tokens=*/false); + EXPECT_EQ(decoded, "hello world"); +} + +TEST_F(RwkvTokenizerTest, CloneUsesModelDirWithRelativeVocabFile) { + RwkvTokenizer tokenizer(test_dir_.string(), MakeRwkvArgs()); + + auto cloned = tokenizer.clone(); + ASSERT_NE(cloned, nullptr); + EXPECT_EQ(cloned->vocab_size(), tokenizer.vocab_size()); + + std::vector ids; + ASSERT_TRUE(cloned->encode("test", &ids, /*add_special_tokens=*/false)); + EXPECT_EQ(ids, std::vector({4})); +} + +TEST_F(RwkvTokenizerTest, AddBosToken) { + TokenizerArgs args = MakeRwkvArgs(); + args.add_bos_token() = true; + args.bos_token() = "<|bos|>"; + args.add_eos_token() = false; + + RwkvTokenizer tokenizer(test_dir_.string(), args); + + std::vector ids; + ASSERT_TRUE( + tokenizer.encode("hello world", &ids, /*add_special_tokens=*/true)); + ASSERT_FALSE(ids.empty()); + EXPECT_EQ(ids.front(), 0); + EXPECT_GT(ids.size(), 1U); +} + +TEST_F(RwkvTokenizerTest, AddEosToken) { + TokenizerArgs args = MakeRwkvArgs(); + args.add_bos_token() = false; + args.add_eos_token() = true; + args.eos_token() = "<|eos|>"; + + RwkvTokenizer tokenizer(test_dir_.string(), args); + + std::vector ids; + ASSERT_TRUE( + tokenizer.encode("hello world", &ids, /*add_special_tokens=*/true)); + ASSERT_FALSE(ids.empty()); + EXPECT_EQ(ids.back(), 1); + EXPECT_GT(ids.size(), 1U); +} + +TEST_F(RwkvTokenizerTest, DecodeSkipSpecialTokens) { + TokenizerArgs args = MakeRwkvArgs(); + args.bos_token() = "<|bos|>"; + args.eos_token() = "<|eos|>"; + + RwkvTokenizer tokenizer(test_dir_.string(), args); + + const std::vector ids = {0, 2, 3, 1}; + const std::string decoded = + tokenizer.decode(ids, /*skip_special_tokens=*/true); + EXPECT_EQ(decoded, "hello world"); +} + +TEST_F(RwkvTokenizerTest, TokenToIdAndIdToToken) { + RwkvTokenizer tokenizer(test_dir_.string(), MakeRwkvArgs()); + + const std::optional hello_id = tokenizer.token_to_id("hello"); + ASSERT_TRUE(hello_id.has_value()); + EXPECT_EQ(hello_id.value(), 2); + EXPECT_EQ(tokenizer.id_to_token(hello_id.value()), "hello"); +} + +TEST_F(RwkvTokenizerTest, EncodeFailsOnUnknownByte) { + RwkvTokenizer tokenizer(test_dir_.string(), MakeRwkvArgs()); + + std::vector ids; + EXPECT_FALSE( + tokenizer.encode("hello\x01", &ids, /*add_special_tokens=*/false)); + EXPECT_TRUE(ids.empty()); +} + +} // namespace xllm diff --git a/tests/core/layers/cuda/CMakeLists.txt b/tests/core/layers/cuda/CMakeLists.txt index 863cc2c100..6c81c165ac 100644 --- a/tests/core/layers/cuda/CMakeLists.txt +++ b/tests/core/layers/cuda/CMakeLists.txt @@ -18,3 +18,24 @@ target_link_libraries(xattention_test PRIVATE brpc "$") + +cc_test( + NAME + rwkv7_decoder_layer_test + SRCS + rwkv7_decoder_layer_test.cpp + DEPS + :config + :layers + :common_layers + :kv_cache + :model_context + :state_dict + torch + GTest::gtest_main + glog::glog +) +target_link_libraries(rwkv7_decoder_layer_test + PRIVATE + brpc + "$") diff --git a/tests/core/layers/cuda/rwkv7_decoder_layer_test.cpp b/tests/core/layers/cuda/rwkv7_decoder_layer_test.cpp new file mode 100644 index 0000000000..ff98150741 --- /dev/null +++ b/tests/core/layers/cuda/rwkv7_decoder_layer_test.cpp @@ -0,0 +1,250 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "core/layers/rwkv7_decoder_layer.h" + +#include +#include +#include + +#include +#include +#include +#include + +#include "framework/kv_cache/kv_cache.h" +#include "framework/model/model_args.h" +#include "framework/model/model_input_params.h" +#include "framework/model_context.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" + +namespace xllm { +namespace layer { +namespace { + +constexpr int64_t kHiddenSize = 128; +constexpr int64_t kHeadSize = 64; +constexpr int64_t kNumHeads = 2; +constexpr int64_t kIntermediateSize = 512; +constexpr int64_t kLoraRank = 16; +constexpr int64_t kGateRank = 32; +constexpr int64_t kNumStateSlots = 4; + +torch::Tensor seeded_tensor(const std::string& name, + const std::vector& shape, + const torch::TensorOptions& options) { + const int64_t seed = static_cast(std::hash{}(name)); + return (torch::randn(shape, options) * 0.02f) + static_cast(seed % 7); +} + +ModelArgs make_model_args() { + ModelArgs args; + args.model_type() = "rwkv7"; + args.hidden_size() = kHiddenSize; + args.head_dim() = kHeadSize; + args.n_heads() = kNumHeads; + args.n_kv_heads() = 1; + args.intermediate_size() = kIntermediateSize; + args.layer_norm_eps() = 1e-5f; + args.linear_num_key_heads() = kNumHeads; + args.linear_num_value_heads() = static_cast(kNumHeads); + args.linear_key_head_dim() = static_cast(kHeadSize); + args.linear_value_head_dim() = static_cast(kHeadSize); + args.linear_conv_kernel_dim() = 2; + return args; +} + +StateDict make_layer0_state_dict(const torch::TensorOptions& options) { + std::unordered_map tensors; + const auto dev = options.device(); + const auto dtype = options.dtype(); + + auto param = [&](const std::string& name) { + return seeded_tensor(name, {1, 1, kHiddenSize}, options); + }; + + tensors["ln0.weight"] = + torch::ones({kHiddenSize}, torch::kFloat32).to(dev).to(dtype); + tensors["ln0.bias"] = + torch::zeros({kHiddenSize}, torch::kFloat32).to(dev).to(dtype); + tensors["ln1.weight"] = tensors["ln0.weight"]; + tensors["ln1.bias"] = tensors["ln0.bias"]; + tensors["ln2.weight"] = tensors["ln0.weight"]; + tensors["ln2.bias"] = tensors["ln0.bias"]; + + tensors["att.x_r"] = param("att.x_r"); + tensors["att.x_w"] = param("att.x_w"); + tensors["att.x_k"] = param("att.x_k"); + tensors["att.x_v"] = param("att.x_v"); + tensors["att.x_a"] = param("att.x_a"); + tensors["att.x_g"] = param("att.x_g"); + tensors["att.w0"] = param("att.w0"); + tensors["att.a0"] = param("att.a0"); + tensors["att.v0"] = param("att.v0"); + tensors["att.k_k"] = torch::ones({1, 1, kHiddenSize}, options) * 0.5f; + tensors["att.k_a"] = torch::ones({1, 1, kHiddenSize}, options) * 0.5f; + tensors["att.r_k"] = + seeded_tensor("att.r_k", {kNumHeads, kHeadSize}, options); + + tensors["att.w1"] = + seeded_tensor("att.w1", {kHiddenSize, kLoraRank}, options); + tensors["att.w2"] = + seeded_tensor("att.w2", {kLoraRank, kHiddenSize}, options); + tensors["att.a1"] = + seeded_tensor("att.a1", {kHiddenSize, kLoraRank}, options); + tensors["att.a2"] = + seeded_tensor("att.a2", {kLoraRank, kHiddenSize}, options); + tensors["att.g1"] = + seeded_tensor("att.g1", {kHiddenSize, kGateRank}, options); + tensors["att.g2"] = + seeded_tensor("att.g2", {kGateRank, kHiddenSize}, options); + + tensors["att.receptance.weight"] = seeded_tensor( + "att.receptance.weight", {kHiddenSize, kHiddenSize}, options); + tensors["att.key.weight"] = + seeded_tensor("att.key.weight", {kHiddenSize, kHiddenSize}, options); + tensors["att.value.weight"] = + seeded_tensor("att.value.weight", {kHiddenSize, kHiddenSize}, options); + tensors["att.output.weight"] = + seeded_tensor("att.output.weight", {kHiddenSize, kHiddenSize}, options); + tensors["att.ln_x.weight"] = + torch::ones({kHiddenSize}, torch::kFloat32).to(dev).to(dtype); + tensors["att.ln_x.bias"] = + torch::zeros({kHiddenSize}, torch::kFloat32).to(dev).to(dtype); + + tensors["ffn.x_k"] = param("ffn.x_k"); + tensors["ffn.key.weight"] = seeded_tensor( + "ffn.key.weight", {kIntermediateSize, kHiddenSize}, options); + tensors["ffn.value.weight"] = seeded_tensor( + "ffn.value.weight", {kHiddenSize, kIntermediateSize}, options); + + return StateDict(std::move(tensors)); +} + +KVCache make_kv_cache(const torch::Device& device) { + auto conv_cache = torch::zeros( + {kNumStateSlots, 1, 3 * kHiddenSize}, + torch::TensorOptions().dtype(torch::kFloat16).device(device)); + auto ssm_cache = torch::zeros( + {kNumStateSlots, kNumHeads, kHeadSize, kHeadSize}, + torch::TensorOptions().dtype(torch::kFloat32).device(device)); + return KVCache(LinearAttentionKVCacheTensors{conv_cache, ssm_cache}); +} + +ModelInputParams make_input_params(int32_t num_sequences, + const std::vector& q_lens, + const torch::Device& device) { + ModelInputParams params; + params.meta.num_sequences = num_sequences; + + std::vector q_cu_seq_lens; + q_cu_seq_lens.reserve(static_cast(q_lens.size()) + 1); + q_cu_seq_lens.push_back(0); + for (int32_t len : q_lens) { + q_cu_seq_lens.push_back(q_cu_seq_lens.back() + len); + } + params.attention.host.q_seq_lens = q_cu_seq_lens; + + std::vector state_indices; + state_indices.reserve(static_cast(num_sequences)); + for (int32_t i = 0; i < num_sequences; ++i) { + state_indices.push_back(i); + } + params.embedding.linear_state_indices = torch::tensor( + state_indices, torch::TensorOptions().dtype(torch::kLong).device(device)); + return params; +} + +} // namespace + +class RWKV7DecoderLayerTest : public ::testing::Test { + protected: + void SetUp() override { + if (!torch::cuda::is_available()) { + GTEST_SKIP() << "CUDA is not available."; + } + device_ = torch::Device(torch::kCUDA, 0); + options_ = torch::TensorOptions().dtype(torch::kFloat16).device(device_); + parallel_args_ = ParallelArgs(1, 1, nullptr); + context_ = + ModelContext(parallel_args_, make_model_args(), QuantArgs(), options_); + layer_ = RWKV7DecoderLayer(context_, /*layer_id=*/0); + layer_->load_state_dict(make_layer0_state_dict(options_)); + kv_cache_ = make_kv_cache(device_); + } + + torch::Device device_{torch::kCPU}; + torch::TensorOptions options_; + ParallelArgs parallel_args_{1, 1, nullptr}; + ModelContext context_{parallel_args_, + make_model_args(), + QuantArgs(), + torch::TensorOptions()}; + RWKV7DecoderLayer layer_{nullptr}; + KVCache kv_cache_; +}; + +TEST_F(RWKV7DecoderLayerTest, ForwardProducesFiniteOutput) { + const int32_t seq_len = 4; + auto input = seeded_tensor("input", {seq_len, kHiddenSize}, options_); + auto params = make_input_params(/*num_sequences=*/1, {seq_len}, device_); + + torch::Tensor v_first; + auto output = + layer_->forward(input, kv_cache_, params, /*layer_id=*/0, v_first); + + EXPECT_EQ(output.sizes(), torch::IntArrayRef({seq_len, kHiddenSize})); + EXPECT_TRUE(torch::isfinite(output).all().item()); + EXPECT_TRUE(v_first.defined()); + EXPECT_EQ(v_first.sizes(), torch::IntArrayRef({seq_len, kHiddenSize})); +} + +TEST_F(RWKV7DecoderLayerTest, ForwardUpdatesLinearAttentionState) { + const int32_t seq_len = 2; + auto input = seeded_tensor("state_input", {seq_len, kHiddenSize}, options_); + auto params = make_input_params(/*num_sequences=*/1, {seq_len}, device_); + + auto conv_before = kv_cache_.get_conv_cache().clone(); + auto ssm_before = kv_cache_.get_ssm_cache().clone(); + + torch::Tensor v_first; + auto output = + layer_->forward(input, kv_cache_, params, /*layer_id=*/0, v_first); + ASSERT_TRUE(output.defined()); + + auto conv_after = kv_cache_.get_conv_cache(); + auto ssm_after = kv_cache_.get_ssm_cache(); + EXPECT_FALSE(torch::allclose(conv_before, conv_after)); + EXPECT_FALSE(torch::allclose(ssm_before, ssm_after)); +} + +TEST_F(RWKV7DecoderLayerTest, ForwardSupportsMultipleSequences) { + const std::vector q_lens = {2, 3}; + const int32_t total_tokens = 5; + auto input = + seeded_tensor("multi_seq_input", {total_tokens, kHiddenSize}, options_); + auto params = make_input_params(/*num_sequences=*/2, q_lens, device_); + + torch::Tensor v_first; + auto output = + layer_->forward(input, kv_cache_, params, /*layer_id=*/0, v_first); + + EXPECT_EQ(output.sizes(), torch::IntArrayRef({total_tokens, kHiddenSize})); + EXPECT_TRUE(torch::isfinite(output).all().item()); +} + +} // namespace layer +} // namespace xllm diff --git a/tools/convert_rwkv7_world.py b/tools/convert_rwkv7_world.py new file mode 100755 index 0000000000..4a833c5b6c --- /dev/null +++ b/tools/convert_rwkv7_world.py @@ -0,0 +1,319 @@ +#!/usr/bin/env python3 +"""Convert BlinkDL / ModelScope RWKV-7 World checkpoints to xLLM model layout. + +xLLM only loads HuggingFace-style model directories. The official RWKV-7-World +release ships native PyTorch checkpoints instead, so this script bridges the +two formats. + +Input format (ModelScope / BlinkDL native) +------------------------------------------ +A single source directory, typically downloaded from ModelScope:: + + rwkv-7-world/ + ├── README.md + ├── configuration.json # {"framework": "Pytorch", "task": "text-generation"} + ├── RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.pth + ├── RWKV-x070-World-0.4B-v2.9-20250107-ctx4096.pth + ├── RWKV-x070-World-1.5B-v3-20250127-ctx4096.pth + └── RWKV-x070-World-2.9B-v3-20250211-ctx4096.pth + +Characteristics of the native format: + +- Weights are stored as raw ``.pth`` state dicts (``emb.weight``, ``blocks.*``). +- There is no ``config.json`` with ``model_type``, layer dims, or context length. +- There is no ``model.safetensors``. +- There is no tokenizer (``tokenizer.json`` / ``tokenizer_config.json``). +- Inference with the official ``rwkv`` pip package reads ``.pth`` directly. + +Output format (xLLM-ready HuggingFace-style layout) +--------------------------------------------------- +One directory per converted checkpoint:: + + rwkv-7-world-0.1b-xllm/ + ├── config.json # model_type=rwkv7, dims inferred from weights + ├── model.safetensors # same tensors as .pth, safetensors format + ├── rwkv_vocab_v20230424.txt # official trie vocab for xLLM rwkv tokenizer + └── tokenizer_config.json # tokenizer_type=rwkv + +``config.json`` fields are inferred from the checkpoint, for example:: + + { + "model_type": "rwkv7", + "torch_dtype": "float16", + "vocab_size": 65536, + "hidden_size": 768, + "num_hidden_layers": 12, + "head_size": 64, + "intermediate_size": 3072, + "num_attention_heads": 12, + "max_position_embeddings": 4096, + ... + } + +Supported sizes +--------------- +- 0.1B: L12, hidden=768 +- 0.4B: L24, hidden=1024 +- 1.5B: L24, hidden=2048 +- 2.9B: L32, hidden=2560 + +Dependencies +------------ +- pip install rwkv safetensors torch + +Usage +----- +:: + + python tools/convert_rwkv7_world.py \\ + --src /path/to/rwkv-7-world \\ + --dst /path/to/output_parent \\ + --sizes all \\ + --dtype float16 + +Alternatively, set ``RWKV7_WORLD_SRC`` and ``RWKV7_WORLD_DST`` instead of +``--src`` / ``--dst``. + +After conversion, start xLLM with:: + + ./build/xllm/core/server/xllm \\ + --model=/path/to/output_parent/rwkv-7-world-0.1b-xllm \\ + --port=9977 --device_id=0 +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import sys +from pathlib import Path + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from scripts.logger import logger + +import torch +from safetensors.torch import save_file + +try: + import rwkv +except ImportError as exc: + logger.error("Missing dependency 'rwkv'. Install with: pip install rwkv") + raise SystemExit(1) from exc + + +ENV_SRC = "RWKV7_WORLD_SRC" +ENV_DST = "RWKV7_WORLD_DST" + +# Known ModelScope filenames for rwkv-7-world. +CHECKPOINTS = [ + { + "pth": "RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.pth", + "suffix": "0.1b-xllm", + "ctx": 4096, + }, + { + "pth": "RWKV-x070-World-0.4B-v2.9-20250107-ctx4096.pth", + "suffix": "0.4b-xllm", + "ctx": 4096, + }, + { + "pth": "RWKV-x070-World-1.5B-v3-20250127-ctx4096.pth", + "suffix": "1.5b-xllm", + "ctx": 4096, + }, + { + "pth": "RWKV-x070-World-2.9B-v3-20250211-ctx4096.pth", + "suffix": "2.9b-xllm", + "ctx": 4096, + }, +] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Convert RWKV-7-World native .pth checkpoints to xLLM " + "HuggingFace-style directories (config.json + model.safetensors " + "+ tokenizer files)." + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "Input: ModelScope dir with RWKV-x070-World-*.pth files only.\n" + "Output: rwkv-7-world--xllm/ per checkpoint." + ), + ) + parser.add_argument( + "--src", + help=( + "Source directory in ModelScope/BlinkDL native format " + f"(required unless {ENV_SRC} is set)" + ), + ) + parser.add_argument( + "--dst", + help=( + "Parent directory for xLLM-ready output folders " + f"(required unless {ENV_DST} is set)" + ), + ) + parser.add_argument( + "--sizes", + default="all", + help="Comma-separated sizes to convert: 0.1b,0.4b,1.5b,2.9b or 'all'", + ) + parser.add_argument( + "--dtype", + default="float16", + choices=("float16", "bfloat16", "float32"), + help="torch_dtype written into config.json", + ) + return parser.parse_args() + + +def infer_arch(state: dict[str, torch.Tensor]) -> dict: + emb = state["emb.weight"] + n_layers = max( + int(key.split(".")[1]) + for key in state + if key.startswith("blocks.") + ) + 1 + head_size = int(state["blocks.0.att.r_k"].shape[-1]) + hidden_size = int(emb.shape[1]) + return { + "vocab_size": int(emb.shape[0]), + "hidden_size": hidden_size, + "num_hidden_layers": n_layers, + "head_size": head_size, + "intermediate_size": int(state["blocks.0.ffn.key.weight"].shape[0]), + "num_attention_heads": hidden_size // head_size, + } + + +def build_config(arch: dict, ctx: int, dtype: str) -> dict: + return { + "model_type": "rwkv7", + "torch_dtype": dtype, + "vocab_size": arch["vocab_size"], + "hidden_size": arch["hidden_size"], + "num_hidden_layers": arch["num_hidden_layers"], + "head_size": arch["head_size"], + "intermediate_size": arch["intermediate_size"], + "num_attention_heads": arch["num_attention_heads"], + "layer_norm_eps": 1e-5, + "max_position_embeddings": ctx, + "bos_token_id": 0, + "eos_token_id": 0, + } + + +def export_tokenizer(out_dir: Path) -> None: + vocab_path = Path(rwkv.__file__).resolve().parent / "rwkv_vocab_v20230424.txt" + if not vocab_path.exists(): + raise FileNotFoundError(f"RWKV vocab not found: {vocab_path}") + + trie_vocab = out_dir / "rwkv_vocab_v20230424.txt" + trie_vocab.write_bytes(vocab_path.read_bytes()) + + with (out_dir / "tokenizer_config.json").open("w", encoding="utf-8") as f: + json.dump( + { + "tokenizer_type": "rwkv", + "tokenizer_class": "RwkvTokenizer", + "vocab_file": "rwkv_vocab_v20230424.txt", + "model_max_length": 4096, + "add_bos_token": False, + "add_eos_token": False, + }, + f, + indent=2, + ) + + # Remove legacy fast-tokenizer export if present. + legacy_tokenizer_json = out_dir / "tokenizer.json" + if legacy_tokenizer_json.exists(): + legacy_tokenizer_json.unlink() + + +def convert_one(src_dir: Path, dst_parent: Path, spec: dict, dtype: str) -> Path: + pth_path = src_dir / spec["pth"] + if not pth_path.exists(): + raise FileNotFoundError(f"Checkpoint not found: {pth_path}") + + out_dir = dst_parent / f"rwkv-7-world-{spec['suffix']}" + out_dir.mkdir(parents=True, exist_ok=True) + + logger.info(f"Converting {pth_path.name} -> {out_dir}") + state = torch.load(pth_path, map_location="cpu", mmap=True, weights_only=True) + arch = infer_arch(state) + config = build_config(arch, spec["ctx"], dtype) + + with (out_dir / "config.json").open("w", encoding="utf-8") as f: + json.dump(config, f, indent=2) + + state = {k: v.contiguous() for k, v in state.items()} + save_file(state, out_dir / "model.safetensors") + export_tokenizer(out_dir) + + logger.info( + f" layers={arch['num_hidden_layers']} hidden={arch['hidden_size']} " + f"vocab={arch['vocab_size']} ctx={spec['ctx']}" + ) + return out_dir + + +def selected_specs(size_arg: str) -> list[dict]: + if size_arg.strip().lower() == "all": + return CHECKPOINTS + + wanted = {s.strip().lower() for s in size_arg.split(",") if s.strip()} + specs = [] + for spec in CHECKPOINTS: + m = re.search(r"(\d+(?:\.\d+)?b)", spec["suffix"], re.I) + label = m.group(1).lower() if m else spec["suffix"] + if label in wanted or spec["suffix"] in wanted: + specs.append(spec) + if not specs: + raise ValueError(f"No checkpoints matched --sizes={size_arg!r}") + return specs + + +def resolve_io_paths(args: argparse.Namespace) -> tuple[Path, Path]: + src = args.src or os.environ.get(ENV_SRC) + dst = args.dst or os.environ.get(ENV_DST) + if not src or not dst: + raise SystemExit( + "Both --src and --dst are required " + f"(or set {ENV_SRC} and {ENV_DST}).\n" + "Example:\n" + " python tools/convert_rwkv7_world.py " + "--src /path/to/rwkv-7-world --dst /path/to/output" + ) + return Path(src).resolve(), Path(dst).resolve() + + +def main() -> int: + args = parse_args() + src_dir, dst_parent = resolve_io_paths(args) + + if not src_dir.is_dir(): + logger.error(f"Source directory not found: {src_dir}") + return 1 + + outputs: list[Path] = [] + for spec in selected_specs(args.sizes): + outputs.append(convert_one(src_dir, dst_parent, spec, args.dtype)) + + logger.info("Done. Start xLLM with one of:") + for out_dir in outputs: + logger.info( + f" ./build/xllm/core/server/xllm " + f"--model={out_dir} --port=9977 --device_id=0" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/xllm/core/framework/tokenizer/CMakeLists.txt b/xllm/core/framework/tokenizer/CMakeLists.txt index 8063c6e99d..a95308529a 100644 --- a/xllm/core/framework/tokenizer/CMakeLists.txt +++ b/xllm/core/framework/tokenizer/CMakeLists.txt @@ -12,6 +12,7 @@ cc_library( tiktoken_tokenizer.h sentencepiece_tokenizer.h fast_tokenizer.h + rwkv_tokenizer.h tokenizer_proxy.h rec_tokenizer.h SRCS @@ -19,6 +20,7 @@ cc_library( tiktoken_tokenizer.cpp sentencepiece_tokenizer.cpp fast_tokenizer.cpp + rwkv_tokenizer.cpp tokenizer_proxy.cpp rec_tokenizer.cpp DEPS diff --git a/xllm/core/framework/tokenizer/rwkv_tokenizer.cpp b/xllm/core/framework/tokenizer/rwkv_tokenizer.cpp new file mode 100644 index 0000000000..523895c0b3 --- /dev/null +++ b/xllm/core/framework/tokenizer/rwkv_tokenizer.cpp @@ -0,0 +1,604 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "rwkv_tokenizer.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace xllm { +namespace { + +struct RwkvTrieNode { + absl::flat_hash_map> children; + std::optional token_id; +}; + +const RwkvTrieNode* find_child(const RwkvTrieNode* node, uint8_t byte) { + const auto it = node->children.find(byte); + if (it == node->children.end()) { + return nullptr; + } + return it->second.get(); +} + +RwkvTrieNode* get_or_create_child(RwkvTrieNode* node, uint8_t byte) { + std::unique_ptr& child = node->children[byte]; + if (child == nullptr) { + child = std::make_unique(); + } + return child.get(); +} + +void utf8_append_codepoint(std::string* out, uint32_t codepoint) { + if (codepoint <= 0x7F) { + out->push_back(static_cast(codepoint)); + } else if (codepoint <= 0x7FF) { + out->push_back(static_cast(0xC0 | ((codepoint >> 6) & 0x1F))); + out->push_back(static_cast(0x80 | (codepoint & 0x3F))); + } else if (codepoint <= 0xFFFF) { + out->push_back(static_cast(0xE0 | ((codepoint >> 12) & 0x0F))); + out->push_back(static_cast(0x80 | ((codepoint >> 6) & 0x3F))); + out->push_back(static_cast(0x80 | (codepoint & 0x3F))); + } else if (codepoint <= 0x10FFFF) { + out->push_back(static_cast(0xF0 | ((codepoint >> 18) & 0x07))); + out->push_back(static_cast(0x80 | ((codepoint >> 12) & 0x3F))); + out->push_back(static_cast(0x80 | ((codepoint >> 6) & 0x3F))); + out->push_back(static_cast(0x80 | (codepoint & 0x3F))); + } +} + +bool parse_hex_digits(const std::string& repr, + size_t pos, + int32_t digits, + uint32_t* value) { + if (value == nullptr || pos + static_cast(digits) > repr.size()) { + return false; + } + uint32_t parsed = 0; + for (int32_t i = 0; i < digits; ++i) { + const char ch = repr[pos + static_cast(i)]; + if (!std::isxdigit(static_cast(ch))) { + return false; + } + parsed = (parsed << 4) + + static_cast(std::isdigit(ch) != 0 + ? ch - '0' + : std::tolower(ch) - 'a' + 10); + } + *value = parsed; + return true; +} + +bool append_codepoint(bool bytes_mode, std::string* out, uint32_t codepoint) { + if (bytes_mode) { + if (codepoint > 0xFF) { + return false; + } + out->push_back(static_cast(codepoint)); + return true; + } + utf8_append_codepoint(out, codepoint); + return true; +} + +bool decode_utf8_codepoint(const std::string& repr, + size_t pos, + uint32_t* codepoint, + size_t* consumed) { + if (codepoint == nullptr || consumed == nullptr || pos >= repr.size()) { + return false; + } + const unsigned char lead = static_cast(repr[pos]); + if ((lead & 0x80) == 0) { + *codepoint = lead; + *consumed = 1; + return true; + } + if ((lead & 0xE0) == 0xC0) { + if (pos + 1 >= repr.size()) { + return false; + } + const unsigned char b1 = static_cast(repr[pos + 1]); + if ((b1 & 0xC0) != 0x80) { + return false; + } + *codepoint = ((lead & 0x1F) << 6) | (b1 & 0x3F); + *consumed = 2; + return true; + } + if ((lead & 0xF0) == 0xE0) { + if (pos + 2 >= repr.size()) { + return false; + } + const unsigned char b1 = static_cast(repr[pos + 1]); + const unsigned char b2 = static_cast(repr[pos + 2]); + if ((b1 & 0xC0) != 0x80 || (b2 & 0xC0) != 0x80) { + return false; + } + *codepoint = ((lead & 0x0F) << 12) | ((b1 & 0x3F) << 6) | (b2 & 0x3F); + *consumed = 3; + return true; + } + if ((lead & 0xF8) == 0xF0) { + if (pos + 3 >= repr.size()) { + return false; + } + const unsigned char b1 = static_cast(repr[pos + 1]); + const unsigned char b2 = static_cast(repr[pos + 2]); + const unsigned char b3 = static_cast(repr[pos + 3]); + if ((b1 & 0xC0) != 0x80 || (b2 & 0xC0) != 0x80 || (b3 & 0xC0) != 0x80) { + return false; + } + *codepoint = ((lead & 0x07) << 18) | ((b1 & 0x3F) << 12) | + ((b2 & 0x3F) << 6) | (b3 & 0x3F); + *consumed = 4; + return true; + } + return false; +} + +bool parse_escape(bool bytes_mode, + const std::string& repr, + size_t* pos, + std::string* out) { + if (pos == nullptr || out == nullptr || *pos >= repr.size()) { + return false; + } + const char esc = repr[(*pos)++]; + switch (esc) { + case '\\': + return append_codepoint(bytes_mode, out, '\\'); + case '\'': + return append_codepoint(bytes_mode, out, '\''); + case '"': + return append_codepoint(bytes_mode, out, '"'); + case 'a': + return append_codepoint(bytes_mode, out, '\a'); + case 'b': + return append_codepoint(bytes_mode, out, '\b'); + case 'f': + return append_codepoint(bytes_mode, out, '\f'); + case 'n': + return append_codepoint(bytes_mode, out, '\n'); + case 'r': + return append_codepoint(bytes_mode, out, '\r'); + case 't': + return append_codepoint(bytes_mode, out, '\t'); + case 'v': + return append_codepoint(bytes_mode, out, '\v'); + case 'x': { + uint32_t value = 0; + if (!parse_hex_digits(repr, *pos, 2, &value)) { + return false; + } + *pos += 2; + return append_codepoint(bytes_mode, out, value); + } + case 'u': { + uint32_t value = 0; + if (!parse_hex_digits(repr, *pos, 4, &value)) { + return false; + } + *pos += 4; + return append_codepoint(bytes_mode, out, value); + } + case 'U': { + uint32_t value = 0; + if (!parse_hex_digits(repr, *pos, 8, &value)) { + return false; + } + *pos += 8; + return append_codepoint(bytes_mode, out, value); + } + default: + return false; + } +} + +// Parse Python str/bytes literals used in rwkv_vocab_v20230424.txt. +// str literals are UTF-8 encoded after parsing, matching official rwkv eval(). +bool parse_python_literal(const std::string& repr, std::string* out) { + if (out == nullptr || repr.empty()) { + return false; + } + out->clear(); + + bool bytes_mode = false; + char quote = 0; + size_t pos = 0; + if (repr.size() >= 2 && repr[0] == 'b' && + (repr[1] == '\'' || repr[1] == '"')) { + bytes_mode = true; + quote = repr[1]; + pos = 2; + } else if (repr[0] == '\'' || repr[0] == '"') { + quote = repr[0]; + pos = 1; + } else { + return false; + } + + while (pos < repr.size()) { + const char ch = repr[pos]; + if (ch == quote) { + return pos + 1 == repr.size(); + } + if (ch == '\\') { + ++pos; + if (!parse_escape(bytes_mode, repr, &pos, out)) { + return false; + } + continue; + } + if (bytes_mode) { + out->push_back(ch); + ++pos; + continue; + } + uint32_t codepoint = 0; + size_t consumed = 0; + if (!decode_utf8_codepoint(repr, pos, &codepoint, &consumed)) { + return false; + } + utf8_append_codepoint(out, codepoint); + pos += consumed; + } + return false; +} + +void add_to_trie(RwkvTrieNode* root, + const std::string& token_bytes, + int32_t token_id) { + RwkvTrieNode* node = root; + for (const unsigned char byte : token_bytes) { + node = get_or_create_child(node, byte); + } + node->token_id = token_id; +} + +} // namespace + +struct RwkvTokenMapHash { + using is_transparent = void; + + size_t operator()(std::string_view key) const { + return absl::Hash{}(key); + } + + size_t operator()(const std::string& key) const { + return absl::Hash{}(std::string_view(key)); + } +}; + +struct RwkvTokenMapEq { + using is_transparent = void; + + bool operator()(std::string_view lhs, std::string_view rhs) const { + return lhs == rhs; + } + + bool operator()(const std::string& lhs, const std::string& rhs) const { + return lhs == rhs; + } + + bool operator()(std::string_view lhs, const std::string& rhs) const { + return lhs == rhs; + } + + bool operator()(const std::string& lhs, std::string_view rhs) const { + return lhs == rhs; + } +}; + +struct RwkvTokenizer::VocabData { + std::unique_ptr root; + absl::flat_hash_map idx_to_token; + absl::flat_hash_map + token_to_idx; + size_t vocab_size = 0; +}; + +std::shared_ptr RwkvTokenizer::build_vocab_data( + const std::string& vocab_file_path) { + auto data = std::make_shared(); + data->root = std::make_unique(); + + std::ifstream input(vocab_file_path); + CHECK(input) << "Failed to open RWKV vocab file: " << vocab_file_path; + + int32_t skipped_lines = 0; + std::string line; + while (std::getline(input, line)) { + if (line.empty()) { + continue; + } + + const size_t first_space = line.find(' '); + const size_t last_space = line.rfind(' '); + if (first_space == std::string::npos || last_space == first_space) { + ++skipped_lines; + continue; + } + + int32_t token_id = 0; + try { + token_id = std::stoi(line.substr(0, first_space)); + } catch (const std::exception&) { + ++skipped_lines; + continue; + } + + const std::string repr = + line.substr(first_space + 1, last_space - first_space - 1); + int32_t expected_len = 0; + try { + expected_len = std::stoi(line.substr(last_space + 1)); + } catch (const std::exception&) { + ++skipped_lines; + continue; + } + + std::string token_bytes; + if (!parse_python_literal(repr, &token_bytes) || + static_cast(token_bytes.size()) != expected_len) { + ++skipped_lines; + continue; + } + + data->idx_to_token[token_id] = token_bytes; + data->token_to_idx[token_bytes] = token_id; + add_to_trie(data->root.get(), token_bytes, token_id); + data->vocab_size = + std::max(data->vocab_size, static_cast(token_id) + 1); + } + + CHECK(!data->idx_to_token.empty()) + << "RWKV vocab file is empty: " << vocab_file_path; + if (skipped_lines > 0) { + LOG(WARNING) << "RWKV vocab skipped " << skipped_lines + << " invalid lines from " << vocab_file_path; + } + LOG(INFO) << "Loaded RWKV trie tokenizer from " << vocab_file_path + << ", entries=" << data->idx_to_token.size() + << ", vocab_size=" << data->vocab_size; + return data; +} + +std::shared_ptr RwkvTokenizer::load_vocab_data( + const std::string& vocab_file_path) { + static std::mutex cache_mutex; + static absl::flat_hash_map> + cache; + + { + std::lock_guard lock(cache_mutex); + const auto it = cache.find(vocab_file_path); + if (it != cache.end()) { + if (auto cached = it->second.lock()) { + return cached; + } + cache.erase(it); + } + } + + auto data = RwkvTokenizer::build_vocab_data(vocab_file_path); + { + std::lock_guard lock(cache_mutex); + cache[vocab_file_path] = data; + } + return data; +} + +RwkvTokenizer::RwkvTokenizer(const std::string_view& dir_path, + const TokenizerArgs& args) + : dir_path_(dir_path), args_(args) { + const std::string vocab_file_path = + dir_path_.empty() ? args_.vocab_file() + : absl::StrCat(dir_path_, "/", args_.vocab_file()); + vocab_ = load_vocab_data(vocab_file_path); +} + +std::unique_ptr RwkvTokenizer::clone() const { + return std::make_unique(dir_path_, args_); +} + +bool RwkvTokenizer::encode_bytes(const std::string_view& text, + std::vector* ids) const { + if (ids == nullptr) { + return false; + } + ids->clear(); + + size_t index = 0; + while (index < text.size()) { + const RwkvTrieNode* node = vocab_->root.get(); + size_t best_end = index; + std::optional best_id; + + size_t cursor = index; + while (cursor < text.size()) { + const uint8_t byte = static_cast(text[cursor]); + const RwkvTrieNode* child = find_child(node, byte); + if (child == nullptr) { + break; + } + node = child; + ++cursor; + if (node->token_id.has_value()) { + best_end = cursor; + best_id = node->token_id; + } + } + + if (!best_id.has_value()) { + LOG(ERROR) << "Failed to tokenize RWKV text at byte offset " << index; + ids->clear(); + return false; + } + ids->push_back(best_id.value()); + index = best_end; + } + return true; +} + +bool RwkvTokenizer::encode(const std::string_view& text, + std::vector* ids, + bool add_special_tokens) const { + if (ids == nullptr) { + return false; + } + + ids->clear(); + if (!encode_bytes(text, ids)) { + return false; + } + if (add_special_tokens && args_.add_bos_token() && + !args_.bos_token().empty()) { + const auto bos_id = token_to_id(args_.bos_token()); + if (bos_id.has_value()) { + ids->insert(ids->begin(), bos_id.value()); + } + } + if (add_special_tokens && args_.add_eos_token() && + !args_.eos_token().empty()) { + const auto eos_id = token_to_id(args_.eos_token()); + if (eos_id.has_value()) { + ids->push_back(eos_id.value()); + } + } + return true; +} + +std::string RwkvTokenizer::decode(const Slice& ids, + bool skip_special_tokens) const { + std::optional bos_id; + std::optional eos_id; + std::optional pad_id; + if (skip_special_tokens) { + if (!args_.bos_token().empty()) { + bos_id = token_to_id(args_.bos_token()); + } + if (!args_.eos_token().empty()) { + eos_id = token_to_id(args_.eos_token()); + } + if (!args_.pad_token().empty()) { + pad_id = token_to_id(args_.pad_token()); + } + } + + std::string bytes; + bytes.reserve(ids.size()); + for (const int32_t id : ids) { + if (skip_special_tokens) { + if (bos_id.has_value() && id == bos_id.value()) { + continue; + } + if (eos_id.has_value() && id == eos_id.value()) { + continue; + } + if (pad_id.has_value() && id == pad_id.value()) { + continue; + } + } + + const auto it = vocab_->idx_to_token.find(id); + if (it == vocab_->idx_to_token.end()) { + LOG(WARNING) << "Unknown RWKV token id during decode: " << id; + continue; + } + bytes.append(it->second); + } + + // Match official RWKV tokenizer behavior: replace invalid UTF-8 with U+FFFD. + std::string output; + output.reserve(bytes.size()); + size_t index = 0; + while (index < bytes.size()) { + const unsigned char byte = static_cast(bytes[index]); + if (byte < 0x80) { + output.push_back(static_cast(byte)); + ++index; + continue; + } + + size_t seq_len = 1; + if ((byte & 0xE0) == 0xC0) { + seq_len = 2; + } else if ((byte & 0xF0) == 0xE0) { + seq_len = 3; + } else if ((byte & 0xF8) == 0xF0) { + seq_len = 4; + } else { + output.append("\xEF\xBF\xBD"); + ++index; + continue; + } + + if (index + seq_len > bytes.size()) { + output.append("\xEF\xBF\xBD"); + break; + } + + bool valid = true; + for (size_t i = 1; i < seq_len; ++i) { + const unsigned char next = static_cast(bytes[index + i]); + if ((next & 0xC0) != 0x80) { + valid = false; + break; + } + } + + if (!valid) { + output.append("\xEF\xBF\xBD"); + ++index; + continue; + } + + output.append(bytes.data() + index, seq_len); + index += seq_len; + } + return output; +} + +std::optional RwkvTokenizer::token_to_id( + const std::string_view& token) const { + const auto it = vocab_->token_to_idx.find(token); + if (it == vocab_->token_to_idx.end()) { + return std::nullopt; + } + return it->second; +} + +std::string RwkvTokenizer::id_to_token(int32_t id) const { + const auto it = vocab_->idx_to_token.find(id); + if (it == vocab_->idx_to_token.end()) { + return ""; + } + return it->second; +} + +size_t RwkvTokenizer::vocab_size() const { return vocab_->vocab_size; } + +} // namespace xllm diff --git a/xllm/core/framework/tokenizer/rwkv_tokenizer.h b/xllm/core/framework/tokenizer/rwkv_tokenizer.h new file mode 100644 index 0000000000..939c0bd7e3 --- /dev/null +++ b/xllm/core/framework/tokenizer/rwkv_tokenizer.h @@ -0,0 +1,66 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include +#include +#include + +#include "tokenizer.h" +#include "tokenizer_args.h" + +namespace xllm { + +// Trie tokenizer used by RWKV-5+ World models (rwkv_vocab_v20230424.txt). +class RwkvTokenizer final : public Tokenizer { + public: + RwkvTokenizer(const std::string_view& dir_path, const TokenizerArgs& args); + + bool encode(const std::string_view& text, + std::vector* ids, + bool add_special_tokens = true) const override; + + std::string decode(const Slice& ids, + bool skip_special_tokens) const override; + + std::optional token_to_id( + const std::string_view& token) const override; + + std::string id_to_token(int32_t id) const override; + + size_t vocab_size() const override; + + std::unique_ptr clone() const override; + + private: + struct VocabData; + + static std::shared_ptr build_vocab_data( + const std::string& vocab_file_path); + + static std::shared_ptr load_vocab_data( + const std::string& vocab_file_path); + + bool encode_bytes(const std::string_view& text, + std::vector* ids) const; + + std::shared_ptr vocab_; + std::string dir_path_; + TokenizerArgs args_; +}; + +} // namespace xllm diff --git a/xllm/core/framework/tokenizer/tokenizer_args.h b/xllm/core/framework/tokenizer/tokenizer_args.h index 0707cae4a6..6251a6482e 100644 --- a/xllm/core/framework/tokenizer/tokenizer_args.h +++ b/xllm/core/framework/tokenizer/tokenizer_args.h @@ -30,8 +30,8 @@ namespace xllm { using SpecialToken = std::pair; struct TokenizerArgs { - // Type of tokenizer to use. valid values are "fast", "sentencepiece" and - // "tiktoken". + // Type of tokenizer to use. valid values are "fast", "sentencepiece", + // "tiktoken", and "rwkv". PROPERTY(std::string, tokenizer_type) = "sentencepiece"; // Vocab file name. diff --git a/xllm/core/framework/tokenizer/tokenizer_factory.cpp b/xllm/core/framework/tokenizer/tokenizer_factory.cpp index 06a1e6fb9d..ebdabacf6f 100644 --- a/xllm/core/framework/tokenizer/tokenizer_factory.cpp +++ b/xllm/core/framework/tokenizer/tokenizer_factory.cpp @@ -24,7 +24,11 @@ std::unique_ptr TokenizerFactory::create_tokenizer( TokenizerArgs tokenizer_args, bool proxy) { std::unique_ptr tokenizer; - if (tokenizer_args.tokenizer_type() == "fast") { + if (tokenizer_args.tokenizer_type() == "rwkv") { + LOG(INFO) << "Create RWKV trie tokenizer."; + tokenizer = + std::make_unique(model_weights_path, tokenizer_args); + } else if (tokenizer_args.tokenizer_type() == "fast") { // 1. fast tokenizer LOG(INFO) << "Create fast tokenizer."; tokenizer = std::make_unique(tokenizer_args); diff --git a/xllm/core/framework/tokenizer/tokenizer_factory.h b/xllm/core/framework/tokenizer/tokenizer_factory.h index 76f7ce9d5a..bf4d2f7894 100644 --- a/xllm/core/framework/tokenizer/tokenizer_factory.h +++ b/xllm/core/framework/tokenizer/tokenizer_factory.h @@ -17,6 +17,7 @@ limitations under the License. #include "fast_tokenizer.h" #include "rec_tokenizer.h" +#include "rwkv_tokenizer.h" #include "sentencepiece_tokenizer.h" #include "tiktoken_tokenizer.h" #include "tokenizer_args.h" diff --git a/xllm/core/layers/CMakeLists.txt b/xllm/core/layers/CMakeLists.txt index 7bcb90fc97..7fb429b1c4 100644 --- a/xllm/core/layers/CMakeLists.txt +++ b/xllm/core/layers/CMakeLists.txt @@ -47,6 +47,7 @@ cc_library( qwen3_vision_layer.h qwen3_decoder_layer.h qwen3_moe_decoder_layer.h + $<$:rwkv7_decoder_layer.h> SRCS oxygen_vision_layer.cpp $<$:deepseek_v4_decoder_layer.cpp> @@ -55,6 +56,7 @@ cc_library( qwen2_5_vision_layer.cpp qwen3_vision_layer.cpp qwen3_moe_decoder_layer.cpp + $<$:rwkv7_decoder_layer.cpp> DEPS $<$:npu_layers> $<$:npu_torch_layers> diff --git a/xllm/core/layers/rwkv7_decoder_layer.cpp b/xllm/core/layers/rwkv7_decoder_layer.cpp new file mode 100644 index 0000000000..cbfe08250f --- /dev/null +++ b/xllm/core/layers/rwkv7_decoder_layer.cpp @@ -0,0 +1,614 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "core/layers/rwkv7_decoder_layer.h" + +#include +#include + +#include +#include +#include +#include +#include + +#include "framework/model/model_args.h" +#include "framework/model/model_input_params.h" +#include "framework/model_context.h" +#include "framework/state_dict/state_dict.h" + +namespace xllm { +namespace layer { + +namespace { + +// ---- Weight-loading helpers ------------------------------------------------ + +// Load a tensor from state_dict and place it on the same device/dtype as ref. +// Returns an undefined tensor if the key is absent. +torch::Tensor load_to(const StateDict& sd, + const std::string& name, + const torch::Device& dev, + torch::ScalarType dtype) { + torch::Tensor t = sd.get_tensor(name); + if (!t.defined()) { + return torch::Tensor(); + } + return t.to(dev).to(dtype); +} + +// Load weight (and optionally bias) for a torch::nn::Linear layer. +void load_linear(torch::nn::Linear& linear, const StateDict& state_dict) { + torch::Tensor w = state_dict.get_tensor("weight"); + if (w.defined()) { + linear->weight = + w.to(linear->weight.device()).to(linear->weight.scalar_type()); + } + if (linear->options.bias()) { + torch::Tensor b = state_dict.get_tensor("bias"); + if (b.defined()) { + linear->bias = b.to(linear->bias.device()).to(linear->bias.scalar_type()); + } + } +} + +// Load weight and bias for a torch::nn::LayerNorm layer. +void load_layer_norm(torch::nn::LayerNorm& ln, const StateDict& state_dict) { + torch::Tensor w = state_dict.get_tensor("weight"); + if (w.defined()) { + ln->weight = w.to(ln->weight.device()).to(ln->weight.scalar_type()); + } + torch::Tensor b = state_dict.get_tensor("bias"); + if (b.defined()) { + ln->bias = b.to(ln->bias.device()).to(ln->bias.scalar_type()); + } +} + +// Load weight and bias for a torch::nn::GroupNorm layer. +void load_group_norm(torch::nn::GroupNorm& gn, const StateDict& state_dict) { + torch::Tensor w = state_dict.get_tensor("weight"); + if (w.defined()) { + gn->weight = w.to(gn->weight.device()).to(gn->weight.scalar_type()); + } + torch::Tensor b = state_dict.get_tensor("bias"); + if (b.defined()) { + gn->bias = b.to(gn->bias.device()).to(gn->bias.scalar_type()); + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// RWKV7TimeMixImpl +// --------------------------------------------------------------------------- + +RWKV7TimeMixImpl::RWKV7TimeMixImpl(const ModelContext& context, + int32_t /*layer_id*/) { + const ModelArgs& args = context.get_model_args(); + hidden_size_ = args.hidden_size(); + head_size_ = args.head_dim(); // RWKV-7 head_size_a + n_heads_ = hidden_size_ / head_size_; + + const torch::TensorOptions opts = context.get_tensor_options(); + + // Token-shift mixing scalars [1, 1, C] — registered parameters so that + // model->to(device) moves them automatically. + x_r_ = register_parameter("x_r", torch::zeros({1, 1, hidden_size_}, opts)); + x_w_ = register_parameter("x_w", torch::zeros({1, 1, hidden_size_}, opts)); + x_k_ = register_parameter("x_k", torch::zeros({1, 1, hidden_size_}, opts)); + x_v_ = register_parameter("x_v", torch::zeros({1, 1, hidden_size_}, opts)); + x_a_ = register_parameter("x_a", torch::zeros({1, 1, hidden_size_}, opts)); + x_g_ = register_parameter("x_g", torch::zeros({1, 1, hidden_size_}, opts)); + + // Decay / a-gate / v-first base scalars [1, 1, C] + w0_ = register_parameter("w0", torch::zeros({1, 1, hidden_size_}, opts)); + a0_ = register_parameter("a0", torch::zeros({1, 1, hidden_size_}, opts)); + v0_ = register_parameter("v0", torch::zeros({1, 1, hidden_size_}, opts)); + + // Key modifiers [1, 1, C] + k_k_ = register_parameter("k_k", torch::ones({1, 1, hidden_size_}, opts)); + k_a_ = register_parameter("k_a", torch::ones({1, 1, hidden_size_}, opts)); + + // Receptance-key interaction [n_heads_, head_size_] + r_k_ = register_parameter("r_k", torch::zeros({n_heads_, head_size_}, opts)); + + // Linear projections: C → C, no bias (BlinkDL checkpoint convention) + receptance_ = register_module( + "receptance", + torch::nn::Linear( + torch::nn::LinearOptions(hidden_size_, hidden_size_).bias(false))); + key_ = register_module( + "key", + torch::nn::Linear( + torch::nn::LinearOptions(hidden_size_, hidden_size_).bias(false))); + value_ = register_module( + "value", + torch::nn::Linear( + torch::nn::LinearOptions(hidden_size_, hidden_size_).bias(false))); + output_ = register_module( + "output", + torch::nn::Linear( + torch::nn::LinearOptions(hidden_size_, hidden_size_).bias(false))); + + // GroupNorm: n_heads_ groups, hidden_size_ channels, eps = 64e-5 (RWKV-7 + // spec) + ln_x_ = register_module( + "ln_x", + torch::nn::GroupNorm(torch::nn::GroupNormOptions( + static_cast(n_heads_), hidden_size_) + .eps(64e-5))); + + receptance_->to(opts.device()); + receptance_->to(opts.dtype().toScalarType()); + key_->to(opts.device()); + key_->to(opts.dtype().toScalarType()); + value_->to(opts.device()); + value_->to(opts.dtype().toScalarType()); + output_->to(opts.device()); + output_->to(opts.dtype().toScalarType()); + ln_x_->to(opts.device()); + ln_x_->to(opts.dtype().toScalarType()); + + // LoRA weights (w1/w2, a1/a2, v1/v2, g1/g2): shapes vary per checkpoint + // rank. They are NOT registered parameters — loaded as plain tensors in + // load_state_dict() and placed on the same device as x_r_. +} + +void RWKV7TimeMixImpl::load_state_dict(const StateDict& state_dict) { + const torch::Device dev = x_r_.device(); + const torch::ScalarType dtype = x_r_.scalar_type(); + + // Registered shift-mix scalars and key modifiers + const std::vector> param_map = { + {"x_r", &x_r_}, + {"x_w", &x_w_}, + {"x_k", &x_k_}, + {"x_v", &x_v_}, + {"x_a", &x_a_}, + {"x_g", &x_g_}, + {"w0", &w0_}, + {"a0", &a0_}, + {"v0", &v0_}, + {"k_k", &k_k_}, + {"k_a", &k_a_}, + {"r_k", &r_k_}, + }; + for (const auto& [name, ptr] : param_map) { + torch::Tensor t = state_dict.get_tensor(name); + if (t.defined()) { + *ptr = t.to(dev).to(dtype); + } + } + + // LoRA tensors: infer rank from checkpoint tensor shapes + w1_ = load_to(state_dict, "w1", dev, dtype); + w2_ = load_to(state_dict, "w2", dev, dtype); + a1_ = load_to(state_dict, "a1", dev, dtype); + a2_ = load_to(state_dict, "a2", dev, dtype); + v1_ = load_to(state_dict, "v1", dev, dtype); + v2_ = load_to(state_dict, "v2", dev, dtype); + g1_ = load_to(state_dict, "g1", dev, dtype); + g2_ = load_to(state_dict, "g2", dev, dtype); + + // Sub-modules: load weight tensors directly + load_linear(receptance_, state_dict.get_dict_with_prefix("receptance.")); + load_linear(key_, state_dict.get_dict_with_prefix("key.")); + load_linear(value_, state_dict.get_dict_with_prefix("value.")); + load_linear(output_, state_dict.get_dict_with_prefix("output.")); + load_group_norm(ln_x_, state_dict.get_dict_with_prefix("ln_x.")); +} + +torch::Tensor RWKV7TimeMixImpl::compute_decay(const torch::Tensor& xw) const { + // xw: [T, C] + // w = log_sigmoid(w0 + tanh(xw @ w1) @ w2) - 0.5 + // ≡ -softplus(-(w0 + tanh(xw @ w1) @ w2)) - 0.5 + // decay = exp(-exp(w)) ∈ (0, 1) + CHECK(w1_.defined() && w2_.defined()) + << "RWKV-7 decay LoRA weights (w1, w2) not loaded"; + torch::Tensor lora = torch::tanh(xw.matmul(w1_)).matmul(w2_); // [T, C] + torch::Tensor w_raw = w0_.view({hidden_size_}) + lora; // [T, C] + torch::Tensor w = torch::log_sigmoid(w_raw) - 0.5f; // [T, C] + torch::Tensor decay = torch::exp(-torch::exp(w)); // [T, C] + return decay.view({-1LL, n_heads_, head_size_}); // [T, H, N] +} + +std::pair RWKV7TimeMixImpl::rwkv7_recurrence( + const torch::Tensor& r, + const torch::Tensor& w, + const torch::Tensor& k, + const torch::Tensor& v, + const torch::Tensor& a, + const torch::Tensor& kk, + const torch::Tensor& state) const { + // r, w, k, v, a, kk: [T, H, N] + // state: [H, N, N] + const int64_t T = r.size(0); + const int64_t H = n_heads_; + const int64_t N = head_size_; + + torch::Tensor out = torch::zeros({T, H * N}, r.options()); + // Float32 for numerical stability of the W-matrix state + torch::Tensor s = state.to(torch::kFloat32); + + for (int64_t t = 0; t < T; ++t) { + // Outer product: vk = v ⊗ k → [H, N, N] + torch::Tensor vt = v[t].to(torch::kFloat32).view({H, N, 1}); + torch::Tensor kt = k[t].to(torch::kFloat32).view({H, 1, N}); + torch::Tensor vk = vt.bmm(kt); + + // In-context learning correction: ab = (-kk) ⊗ (kk * a) → [H, N, N] + torch::Tensor kkt = kk[t].to(torch::kFloat32); // [H, N] + torch::Tensor at = a[t].to(torch::kFloat32); // [H, N] + torch::Tensor ab = (-kkt).view({H, N, 1}).bmm((kkt * at).view({H, 1, N})); + + // State update: s = s * w + s @ ab + vk + torch::Tensor wt = w[t].to(torch::kFloat32).view({H, 1, N}); + s = s * wt + s.bmm(ab) + vk; + + // Read-out: out[t] = (s @ r)^T → [H, N] + torch::Tensor rt = r[t].to(torch::kFloat32).view({H, N, 1}); + out[t] = s.bmm(rt).view({H * N}).to(r.dtype()); + } + + return {out, s.to(state.dtype())}; +} + +std::pair RWKV7TimeMixImpl::forward( + const torch::Tensor& x, + const torch::Tensor& att_x_prev, + torch::Tensor& att_kv, + bool is_layer0, + torch::Tensor& v_first, + const std::vector& seq_lens) { + // x: [total_tokens, C] + // att_x_prev: [num_seqs, C] + // att_kv: [num_seqs, H, N, N] + + const int64_t total_tokens = x.size(0); + const int32_t num_seqs = static_cast(seq_lens.size()); + + torch::Tensor output = torch::zeros_like(x); + torch::Tensor new_att_x_prev = torch::zeros_like(att_x_prev); + + if (is_layer0 && !v_first.defined()) { + v_first = torch::zeros_like(x); + } + + int64_t token_offset = 0; + for (int32_t s = 0; s < num_seqs; ++s) { + const int64_t T = static_cast(seq_lens[static_cast(s)]); + torch::Tensor xs = x.slice(0, token_offset, token_offset + T); // [T, C] + torch::Tensor x_prev_s = att_x_prev[s].unsqueeze(0); // [1, C] + + // Time-shift: xx[t] = x[t-1] - x[t], with x[-1] = att_x_prev + torch::Tensor shifted; + if (T > 1) { + shifted = torch::cat({x_prev_s, xs.slice(0, 0, T - 1)}, 0); + } else { + shifted = x_prev_s; + } + torch::Tensor xx = shifted - xs; // [T, C] + + // Mixed inputs using shift scalars (broadcast [1,1,C] → [T,C]) + auto mix = [&](const torch::Tensor& coeff) { + return xs + xx * coeff.view({hidden_size_}); // [T, C] + }; + torch::Tensor xr = mix(x_r_); + torch::Tensor xw = mix(x_w_); + torch::Tensor xk = mix(x_k_); + torch::Tensor xv = mix(x_v_); + torch::Tensor xa = mix(x_a_); + torch::Tensor xg = mix(x_g_); + + // Linear projections + torch::Tensor r_proj = receptance_(xr); // [T, C] + torch::Tensor k_proj = key_(xk); // [T, C] + torch::Tensor v_proj = value_(xv); // [T, C] + + // Decay [T, H, N] + torch::Tensor w_decay = compute_decay(xw); + + // A-gate: sigmoid(a0 + xa @ a1 @ a2) → [T, C] + CHECK(a1_.defined() && a2_.defined()) + << "RWKV-7 a-gate LoRA weights (a1, a2) not loaded"; + torch::Tensor a_gate = + torch::sigmoid(a0_.view({hidden_size_}) + xa.matmul(a1_).matmul(a2_)); + + // Gate: sigmoid(xg @ g1) @ g2 → [T, C] + CHECK(g1_.defined() && g2_.defined()) + << "RWKV-7 gate LoRA weights (g1, g2) not loaded"; + torch::Tensor gate = torch::sigmoid(xg.matmul(g1_)).matmul(g2_); + + // Head layout reshaping: [T, H, N] + torch::Tensor r_h = r_proj.view({T, n_heads_, head_size_}); + torch::Tensor k_h = k_proj.view({T, n_heads_, head_size_}); + torch::Tensor v_h = v_proj.view({T, n_heads_, head_size_}); + torch::Tensor a_h = a_gate.view({T, n_heads_, head_size_}); + + // Key normalisation (per-head L2 norm) + torch::Tensor kk_h = torch::nn::functional::normalize( + k_h * k_k_.view({1LL, n_heads_, head_size_}), + torch::nn::functional::NormalizeFuncOptions().dim(-1).p(2)); + + // Key modulation with a-gate + torch::Tensor k_mod = + k_h * (1.0f + (a_h - 1.0f) * k_a_.view({1LL, n_heads_, head_size_})); + + // V-first blending + torch::Tensor v_h_blend = v_h; + if (is_layer0) { + v_first.slice(0, token_offset, token_offset + T).copy_(v_proj); + } else { + CHECK(v1_.defined() && v2_.defined()) + << "RWKV-7 v-first LoRA weights (v1, v2) not loaded"; + torch::Tensor vf = v_first.slice(0, token_offset, token_offset + T) + .view({T, n_heads_, head_size_}); + torch::Tensor blend = + torch::sigmoid(v0_.view({hidden_size_}) + xv.matmul(v1_).matmul(v2_)) + .view({T, n_heads_, head_size_}); + v_h_blend = v_h + (vf - v_h) * blend; + } + + // Core RWKV-7 recurrence + auto [seq_out, new_state] = + rwkv7_recurrence(r_h, w_decay, k_mod, v_h_blend, a_h, kk_h, att_kv[s]); + att_kv[s].copy_(new_state); // update W-matrix state in-place + + // GroupNorm over the channel dimension + torch::Tensor normed = ln_x_(seq_out.view({T, hidden_size_})); + + // Receptance-key residual: per-head dot-product correction + torch::Tensor rk_dot = (r_h * k_mod * r_k_.unsqueeze(0)) + .sum(/*dim=*/-1, /*keepdim=*/true); // [T, H, 1] + torch::Tensor rk_res = (rk_dot * v_h_blend).view({T, hidden_size_}); + + // Output projection + torch::Tensor block_out = output_((normed + rk_res) * gate); // [T, C] + + output.slice(0, token_offset, token_offset + T).copy_(block_out); + new_att_x_prev[s].copy_(xs[T - 1]); // last time-mix input for next shift + + token_offset += T; + } + + return {output, new_att_x_prev}; +} + +// --------------------------------------------------------------------------- +// RWKV7ChannelMixImpl +// --------------------------------------------------------------------------- + +RWKV7ChannelMixImpl::RWKV7ChannelMixImpl(const ModelContext& context, + int32_t /*layer_id*/) { + const ModelArgs& args = context.get_model_args(); + const int64_t hidden_size = args.hidden_size(); + const int64_t intermediate_size = args.intermediate_size(); + const torch::TensorOptions opts = context.get_tensor_options(); + + x_k_ = register_parameter("x_k", torch::zeros({1, 1, hidden_size}, opts)); + + key_ = register_module( + "key", + torch::nn::Linear(torch::nn::LinearOptions(hidden_size, intermediate_size) + .bias(false))); + value_ = register_module( + "value", + torch::nn::Linear(torch::nn::LinearOptions(intermediate_size, hidden_size) + .bias(false))); + key_->to(opts.device()); + key_->to(opts.dtype().toScalarType()); + value_->to(opts.device()); + value_->to(opts.dtype().toScalarType()); +} + +void RWKV7ChannelMixImpl::load_state_dict(const StateDict& state_dict) { + torch::Tensor xk_t = state_dict.get_tensor("x_k"); + if (xk_t.defined()) { + x_k_ = xk_t.to(x_k_.device()).to(x_k_.scalar_type()); + } + load_linear(key_, state_dict.get_dict_with_prefix("key.")); + load_linear(value_, state_dict.get_dict_with_prefix("value.")); +} + +std::pair RWKV7ChannelMixImpl::forward( + const torch::Tensor& x, + const torch::Tensor& ffn_x_prev, + const std::vector& seq_lens) { + const int32_t num_seqs = static_cast(seq_lens.size()); + const int64_t C = x.size(-1); + + torch::Tensor output = torch::zeros_like(x); + torch::Tensor new_ffn_x_prev = torch::zeros_like(ffn_x_prev); + + int64_t token_offset = 0; + for (int32_t s = 0; s < num_seqs; ++s) { + const int64_t T = static_cast(seq_lens[static_cast(s)]); + torch::Tensor xs = x.slice(0, token_offset, token_offset + T); + torch::Tensor x_prev_s = ffn_x_prev[s].unsqueeze(0); // [1, C] + + torch::Tensor shifted; + if (T > 1) { + shifted = torch::cat({x_prev_s, xs.slice(0, 0, T - 1)}, 0); + } else { + shifted = x_prev_s; + } + torch::Tensor xx = shifted - xs; + torch::Tensor k = xs + xx * x_k_.view({C}); // [T, C] + + // Gated FFN: relu(key(k)) ^ 2 → value + torch::Tensor k_act = torch::pow(torch::relu(key_(k)), 2.0f); + torch::Tensor out_s = value_(k_act); // [T, C] + + output.slice(0, token_offset, token_offset + T).copy_(out_s); + new_ffn_x_prev[s].copy_(xs[T - 1]); + + token_offset += T; + } + + return {output, new_ffn_x_prev}; +} + +// --------------------------------------------------------------------------- +// RWKV7DecoderLayerImpl +// --------------------------------------------------------------------------- + +RWKV7DecoderLayerImpl::RWKV7DecoderLayerImpl(const ModelContext& context, + int32_t layer_id) + : layer_id_(layer_id), has_ln0_(layer_id == 0) { + const ModelArgs& args = context.get_model_args(); + hidden_size_ = args.hidden_size(); + head_size_ = args.head_dim(); + n_heads_ = hidden_size_ / head_size_; + + const float ln_eps = + args.layer_norm_eps() > 0.0f ? args.layer_norm_eps() : 1e-5f; + const torch::TensorOptions opts = context.get_tensor_options(); + + if (has_ln0_) { + ln0_ = register_module( + "ln0", + torch::nn::LayerNorm( + torch::nn::LayerNormOptions({hidden_size_}).eps(ln_eps))); + ln0_->to(opts.device()); + ln0_->to(opts.dtype().toScalarType()); + } + ln1_ = register_module( + "ln1", + torch::nn::LayerNorm( + torch::nn::LayerNormOptions({hidden_size_}).eps(ln_eps))); + ln2_ = register_module( + "ln2", + torch::nn::LayerNorm( + torch::nn::LayerNormOptions({hidden_size_}).eps(ln_eps))); + ln1_->to(opts.device()); + ln1_->to(opts.dtype().toScalarType()); + ln2_->to(opts.device()); + ln2_->to(opts.dtype().toScalarType()); + + att_ = register_module("att", RWKV7TimeMix(context, layer_id)); + ffn_ = register_module("ffn", RWKV7ChannelMix(context, layer_id)); +} + +void RWKV7DecoderLayerImpl::load_state_dict(const StateDict& state_dict) { + if (has_ln0_) { + load_layer_norm(ln0_, state_dict.get_dict_with_prefix("ln0.")); + } + load_layer_norm(ln1_, state_dict.get_dict_with_prefix("ln1.")); + load_layer_norm(ln2_, state_dict.get_dict_with_prefix("ln2.")); + att_->load_state_dict(state_dict.get_dict_with_prefix("att.")); + ffn_->load_state_dict(state_dict.get_dict_with_prefix("ffn.")); +} + +std::tuple +RWKV7DecoderLayerImpl::read_state(KVCache& kv_cache, + const torch::Tensor& state_indices, + int32_t /*num_seqs*/) const { + torch::Tensor conv = kv_cache.get_conv_cache(); + torch::Tensor ssm = kv_cache.get_ssm_cache(); + + CHECK(conv.defined()) << "conv_cache is not allocated; ensure RWKV-7 uses " + "the linear-attention KV cache."; + CHECK(ssm.defined()) << "ssm_cache is not allocated; ensure RWKV-7 uses " + "the linear-attention KV cache."; + + torch::Tensor sel_conv = conv.index_select(0, state_indices); // [S, 1, 3H] + torch::Tensor sel_ssm = ssm.index_select(0, state_indices); // [S, H, N, N] + + // Unpack shift states from conv_cache[slot, 0, 0:2H] + torch::Tensor flat = sel_conv.squeeze(1); // [S, 3H] + torch::Tensor att_x_prev = + flat.slice(/*dim=*/-1, 0LL, hidden_size_).contiguous(); + torch::Tensor ffn_x_prev = + flat.slice(/*dim=*/-1, hidden_size_, 2LL * hidden_size_).contiguous(); + + return {att_x_prev, ffn_x_prev, sel_ssm.contiguous()}; +} + +void RWKV7DecoderLayerImpl::write_state(KVCache& kv_cache, + const torch::Tensor& state_indices, + const torch::Tensor& att_x_prev, + const torch::Tensor& ffn_x_prev, + const torch::Tensor& att_kv) const { + torch::Tensor conv = kv_cache.get_conv_cache(); + torch::Tensor ssm = kv_cache.get_ssm_cache(); + + const int64_t num_seqs = att_x_prev.size(0); + torch::Tensor new_conv = + torch::zeros({num_seqs, 1LL, 3LL * hidden_size_}, conv.options()); + + // Pack shift states into new_conv[s, 0, 0:H] and [H:2H] + new_conv.squeeze(1).slice(/*dim=*/-1, 0LL, hidden_size_).copy_(att_x_prev); + new_conv.squeeze(1) + .slice(/*dim=*/-1, hidden_size_, 2LL * hidden_size_) + .copy_(ffn_x_prev); + + conv.index_copy_(0, state_indices, new_conv); + ssm.index_copy_(0, state_indices, att_kv); +} + +torch::Tensor RWKV7DecoderLayerImpl::forward( + const torch::Tensor& x, + KVCache& kv_cache, + const ModelInputParams& input_params, + int32_t layer_id, + torch::Tensor& v_first) { + const int32_t num_seqs = input_params.meta.num_sequences; + + std::vector seq_lens(static_cast(num_seqs)); + for (int32_t s = 0; s < num_seqs; ++s) { + seq_lens[static_cast(s)] = input_params.get_q_seq_len(s); + } + + // Use scheduler-provided slot indices when available; otherwise fall back + // to sequential slot assignment [0, 1, …, num_seqs-1]. + torch::Tensor state_indices; + if (input_params.embedding.linear_state_indices.defined()) { + state_indices = input_params.embedding.linear_state_indices; + } else { + state_indices = torch::arange( + num_seqs, + torch::TensorOptions().dtype(torch::kLong).device(x.device())); + } + if (state_indices.scalar_type() != torch::kLong) { + state_indices = state_indices.to(torch::kLong); + } + + auto [att_x_prev, ffn_x_prev, att_kv] = + read_state(kv_cache, state_indices, num_seqs); + + // ln0: applied at block 0 before ln1 to normalise the token embedding + torch::Tensor h = x; + if (has_ln0_) { + h = ln0_(h); + } + + // Time-mix (attention replacement) + torch::Tensor h_att = ln1_(h); + auto [att_out, new_att_x_prev] = att_->forward( + h_att, att_x_prev, att_kv, layer_id == 0, v_first, seq_lens); + h = h + att_out; + + // Channel-mix (FFN replacement) + torch::Tensor h_ffn = ln2_(h); + auto [ffn_out, new_ffn_x_prev] = ffn_->forward(h_ffn, ffn_x_prev, seq_lens); + h = h + ffn_out; + + write_state(kv_cache, state_indices, new_att_x_prev, new_ffn_x_prev, att_kv); + + return h; +} + +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/rwkv7_decoder_layer.h b/xllm/core/layers/rwkv7_decoder_layer.h new file mode 100644 index 0000000000..54789be0d8 --- /dev/null +++ b/xllm/core/layers/rwkv7_decoder_layer.h @@ -0,0 +1,246 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// RWKV-7 "Goose" decoder layer for xLLM. +// +// Implements the RWKV-7 time-mix (attention replacement) and channel-mix +// (FFN replacement) blocks. Architecture reference: +// https://github.com/BlinkDL/RWKV-LM/tree/main/RWKV-v7 +// +// Weight naming convention follows the official BlinkDL checkpoint format: +// blocks.{i}.att.* → time-mix weights +// blocks.{i}.ffn.* → channel-mix weights +// blocks.{i}.ln0/ln1/ln2 → layer norms + +#pragma once + +#include + +#include +#include +#include + +#include "framework/kv_cache/kv_cache.h" +#include "framework/model/model_args.h" +#include "framework/model/model_input_params.h" +#include "framework/model_context.h" +#include "framework/state_dict/state_dict.h" + +namespace xllm { +namespace layer { + +// RWKV-7 Time-Mix block (attention replacement). +// +// Key operations per token: +// 1. Token-shift mixing: xx = prev_token_embed - current_embed +// 2. Compute gating values via LoRA-style low-rank projections +// 3. In-context learning (a-gate) modulates the W-matrix update +// 4. W-matrix recurrence: state = state * decay + state @ ab + vk +// 5. Receptance-key residual connection for output correction +class RWKV7TimeMixImpl : public torch::nn::Module { + public: + RWKV7TimeMixImpl(const ModelContext& context, int32_t layer_id); + + void load_state_dict(const StateDict& state_dict); + + // Forward time-mix for all sequences in the current batch. + // + // x – layer-normed block input: [total_tokens, hidden_size] + // att_x_prev – per-sequence shift state: [num_seqs, hidden_size] + // att_kv – per-sequence W-matrix: [num_seqs, n_heads, head_size, + // head_size] + // Updated in-place during this call. + // is_layer0 – true for block 0 (initialises v_first) + // v_first – cross-layer value state: [total_tokens, hidden_size] + // (in/out) seq_lens – per-sequence query-token counts (host vector) + // + // Returns {time_mix_output [total_tokens, C], + // new_att_x_prev [num_seqs, C]} + std::pair forward( + const torch::Tensor& x, + const torch::Tensor& att_x_prev, + torch::Tensor& att_kv, + bool is_layer0, + torch::Tensor& v_first, + const std::vector& seq_lens); + + private: + // Compute per-token decay exp(-exp(w)) for one sequence. + // xw: [T, C] → output: [T, n_heads_, head_size_] + torch::Tensor compute_decay(const torch::Tensor& xw) const; + + // Core RWKV-7 recurrence kernel (sequential over T) for one sequence. + // r, w, k, v, a, kk: [T, n_heads_, head_size_] + // state: [n_heads_, head_size_, head_size_] + // Returns {output [T, C], new_state [n_heads_, head_size_, head_size_]}. + std::pair rwkv7_recurrence( + const torch::Tensor& r, + const torch::Tensor& w, + const torch::Tensor& k, + const torch::Tensor& v, + const torch::Tensor& a, + const torch::Tensor& kk, + const torch::Tensor& state) const; + + int64_t hidden_size_; + int64_t n_heads_; + int64_t head_size_; + + // Token-shift mixing scalars — registered parameters [1, 1, hidden_size]. + torch::Tensor x_r_, x_w_, x_k_, x_v_, x_a_, x_g_; + + // Decay LoRA: w = log_sigmoid(w0 + xw@w1@w2) - 0.5 + // w0: [1, 1, C], w1: [C, rank], w2: [rank, C] + // Shapes of w1/w2 depend on the checkpoint rank; loaded dynamically. + torch::Tensor w0_; // registered parameter [1, 1, C] + torch::Tensor w1_; // plain tensor, shape from checkpoint + torch::Tensor w2_; // plain tensor, shape from checkpoint + + // A-gate LoRA: a = sigmoid(a0 + xa@a1@a2) + torch::Tensor a0_; // registered parameter [1, 1, C] + torch::Tensor a1_; // plain tensor + torch::Tensor a2_; // plain tensor + + // V-first blend LoRA: blend = sigmoid(v0 + xv@v1@v2) + torch::Tensor v0_; // registered parameter [1, 1, C] + torch::Tensor v1_; // plain tensor + torch::Tensor v2_; // plain tensor + + // Gate LoRA: g = sigmoid(xg@g1) @ g2 + torch::Tensor g1_; // plain tensor + torch::Tensor g2_; // plain tensor + + // Key modifiers — registered parameters [1, 1, hidden_size] + torch::Tensor k_k_; + torch::Tensor k_a_; + + // Receptance-key interaction — registered parameter [n_heads_, head_size_] + torch::Tensor r_k_; + + // Linear projections (no bias to match BlinkDL checkpoint format) + torch::nn::Linear receptance_{nullptr}; + torch::nn::Linear key_{nullptr}; + torch::nn::Linear value_{nullptr}; + torch::nn::Linear output_{nullptr}; + + // Post-attention group norm: GroupNorm(n_heads_, hidden_size_, eps=64e-5) + torch::nn::GroupNorm ln_x_{nullptr}; +}; + +TORCH_MODULE(RWKV7TimeMix); + +// RWKV-7 Channel-Mix block (FFN replacement). +// +// Simple gated MLP with token shift: +// k = x + (prev_token - x) * x_k +// output = value(relu(key(k)) ^ 2) +class RWKV7ChannelMixImpl : public torch::nn::Module { + public: + RWKV7ChannelMixImpl(const ModelContext& context, int32_t layer_id); + + void load_state_dict(const StateDict& state_dict); + + // Forward channel-mix for all sequences in the current batch. + // + // x – layer-normed input: [total_tokens, hidden_size] + // ffn_x_prev – per-sequence shift state: [num_seqs, hidden_size] + // seq_lens – per-sequence query-token counts + // + // Returns {channel_mix_output [total_tokens, C], + // new_ffn_x_prev [num_seqs, C]} + std::pair forward( + const torch::Tensor& x, + const torch::Tensor& ffn_x_prev, + const std::vector& seq_lens); + + private: + // Token-shift mixing scalar — registered parameter [1, 1, hidden_size] + torch::Tensor x_k_; + + // Linear projections: key expands, value contracts (no bias) + torch::nn::Linear key_{nullptr}; // hidden_size → intermediate_size + torch::nn::Linear value_{nullptr}; // intermediate_size → hidden_size +}; + +TORCH_MODULE(RWKV7ChannelMix); + +// Complete RWKV-7 block: (ln0 →) ln1 → time-mix → residual → ln2 → channel-mix +// → residual. +// +// ln0 is applied only in block 0 as a pre-normalisation of the token embedding. +class RWKV7DecoderLayerImpl : public torch::nn::Module { + public: + RWKV7DecoderLayerImpl(const ModelContext& context, int32_t layer_id); + + void load_state_dict(const StateDict& state_dict); + + // Forward one RWKV-7 block. + // + // x – block input: [total_tokens, hidden_size] + // kv_cache – this layer's linear-attention state (conv + ssm tensors) + // input_params – batch metadata (seq_lens, linear_state_ids …) + // layer_id – 0-based layer index (0 triggers ln0 + v_first + // initialisation) v_first – cross-layer value state, updated in-place at + // layer 0 + // + // Returns the block output: [total_tokens, hidden_size]. + torch::Tensor forward(const torch::Tensor& x, + KVCache& kv_cache, + const ModelInputParams& input_params, + int32_t layer_id, + torch::Tensor& v_first); + + private: + // Read per-sequence shift states and W-matrix from the KV cache. + // + // conv_cache layout (dim -1 = 3 * hidden_size): + // [0 : H] → att_x_prev + // [H : 2H] → ffn_x_prev + // [2H : 3H] → unused (artefact of conv_cache shape formula) + // + // ssm_cache: [num_slots, n_heads_, head_size_, head_size_] → W-matrix state. + // + // Returns {att_x_prev [seqs, H], ffn_x_prev [seqs, H], att_kv [seqs, H, N, + // N]}. + std::tuple read_state( + KVCache& kv_cache, + const torch::Tensor& state_indices, + int32_t num_seqs) const; + + // Write updated states back to the KV cache. + void write_state(KVCache& kv_cache, + const torch::Tensor& state_indices, + const torch::Tensor& att_x_prev, + const torch::Tensor& ffn_x_prev, + const torch::Tensor& att_kv) const; + + torch::nn::LayerNorm ln0_{nullptr}; // only registered for block 0 + torch::nn::LayerNorm ln1_{nullptr}; + torch::nn::LayerNorm ln2_{nullptr}; + + RWKV7TimeMix att_{nullptr}; + RWKV7ChannelMix ffn_{nullptr}; + + int64_t hidden_size_; + int64_t n_heads_; + int64_t head_size_; + int32_t layer_id_; + bool has_ln0_; +}; + +TORCH_MODULE(RWKV7DecoderLayer); + +} // namespace layer +} // namespace xllm diff --git a/xllm/models/llm/rwkv7.h b/xllm/models/llm/rwkv7.h new file mode 100644 index 0000000000..994a71949a --- /dev/null +++ b/xllm/models/llm/rwkv7.h @@ -0,0 +1,349 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// RWKV-7 "Goose" causal LM implementation for xLLM inference. +// +// Architecture reference: +// https://github.com/BlinkDL/RWKV-LM/tree/main/RWKV-v7 +// +// Key characteristics that distinguish RWKV-7 from transformer LLMs: +// • No positional encoding (recurrence replaces attention). +// • Fixed-size per-layer recurrent state (W-matrix + shift states). +// • State stored in the linear-attention KV cache (conv + ssm tensors). +// • v_first cross-layer tensor managed in the custom forward loop. +// +// Checkpoint weight layout (BlinkDL / HuggingFace RWKV): +// emb.weight, blocks.{i}.{ln0/ln1/ln2/att/ffn}.*, ln_out.*, head.weight + +#pragma once + +#include + +#include +#include +#include +#include + +#include "core/framework/kv_cache/kv_cache.h" +#include "core/framework/model/model_args.h" +#include "core/framework/model/model_input_params.h" +#include "core/framework/model/model_output.h" +#include "core/framework/model_context.h" +#include "core/framework/model_loader.h" +#include "core/layers/rwkv7_decoder_layer.h" +#include "models/model_registry.h" + +namespace xllm { + +// --------------------------------------------------------------------------- +// RWKV7Model – embedding + N blocks + final layer-norm +// --------------------------------------------------------------------------- + +class RWKV7ModelImpl : public torch::nn::Module { + public: + explicit RWKV7ModelImpl(const ModelContext& context) { + const ModelArgs& args = context.get_model_args(); + const int32_t n_layers = static_cast(args.n_layers()); + const int64_t hidden_size = args.hidden_size(); + const int64_t vocab_size = args.vocab_size(); + const float ln_eps = + args.layer_norm_eps() > 0.0f ? args.layer_norm_eps() : 1e-5f; + const torch::TensorOptions opts = context.get_tensor_options(); + + emb_ = register_module("emb", + torch::nn::Embedding(torch::nn::EmbeddingOptions( + vocab_size, hidden_size))); + emb_->to(opts.device()); + emb_->to(opts.dtype().toScalarType()); + + blocks_module_list_ = register_module("blocks", torch::nn::ModuleList()); + blocks_.reserve(static_cast(n_layers)); + for (int32_t i = 0; i < n_layers; ++i) { + auto block = layer::RWKV7DecoderLayer(context, i); + blocks_.emplace_back(block); + blocks_module_list_->push_back(block); + } + + ln_out_ = register_module( + "ln_out", + torch::nn::LayerNorm( + torch::nn::LayerNormOptions({hidden_size}).eps(ln_eps))); + ln_out_->to(opts.device()); + ln_out_->to(opts.dtype().toScalarType()); + } + + // tokens: [total_tokens] (all sequences packed) + // positions: unused — RWKV-7 has no positional encoding + // kv_caches: one KVCache per layer; each holds conv+ssm state tensors + // + // Returns: ModelOutput with hidden_states [total_tokens, hidden_size] + ModelOutput forward(const torch::Tensor& tokens, + const torch::Tensor& /*positions*/, + std::vector& kv_caches, + const ModelInputParams& input_params) { + torch::Tensor h = emb_(tokens); // [total_tokens, hidden_size] + + // v_first is a cross-layer state initialised by block 0 and passed + // through all subsequent blocks for value-residual blending. + torch::Tensor v_first; + + for (size_t i = 0; i < blocks_.size(); ++i) { + h = blocks_[i]->forward( + h, kv_caches[i], input_params, static_cast(i), v_first); + } + + torch::Tensor hidden_states = ln_out_(h); + return ModelOutput(hidden_states, std::nullopt); + } + + void load_state_dict(const StateDict& state_dict) { + // Embedding: RWKV checkpoints use "emb.weight" + torch::Tensor emb_w = + state_dict.get_dict_with_prefix("emb.").get_tensor("weight"); + if (emb_w.defined()) { + emb_->weight = + emb_w.to(emb_->weight.device()).to(emb_->weight.scalar_type()); + } + + // Decoder blocks + for (size_t i = 0; i < blocks_.size(); ++i) { + const std::string prefix = "blocks." + std::to_string(i) + "."; + blocks_[i]->load_state_dict(state_dict.get_dict_with_prefix(prefix)); + } + + // Final layer norm: load weight and bias directly (xLLM StateDict API) + const StateDict ln_dict = state_dict.get_dict_with_prefix("ln_out."); + torch::Tensor ln_w = ln_dict.get_tensor("weight"); + if (ln_w.defined()) { + ln_out_->weight = + ln_w.to(ln_out_->weight.device()).to(ln_out_->weight.scalar_type()); + } + torch::Tensor ln_b = ln_dict.get_tensor("bias"); + if (ln_b.defined()) { + ln_out_->bias = + ln_b.to(ln_out_->bias.device()).to(ln_out_->bias.scalar_type()); + } + } + + private: + torch::nn::Embedding emb_{nullptr}; + torch::nn::ModuleList blocks_module_list_{nullptr}; + std::vector blocks_; + torch::nn::LayerNorm ln_out_{nullptr}; +}; + +TORCH_MODULE(RWKV7Model); + +// --------------------------------------------------------------------------- +// RWKV7ForCausalLM – adds the unembedding LM head +// --------------------------------------------------------------------------- + +class RWKV7ForCausalLMImpl : public torch::nn::Module { + public: + explicit RWKV7ForCausalLMImpl(const ModelContext& context) { + const ModelArgs& args = context.get_model_args(); + const int64_t hidden_size = args.hidden_size(); + const int64_t vocab_size = args.vocab_size(); + const torch::TensorOptions opts = context.get_tensor_options(); + + model_ = register_module("model", RWKV7Model(context)); + + // LM head: linear, no bias (BlinkDL convention) + head_ = register_module( + "head", + torch::nn::Linear( + torch::nn::LinearOptions(hidden_size, vocab_size).bias(false))); + head_->to(opts.device()); + head_->to(opts.dtype().toScalarType()); + } + + // Forward: compute hidden states for the packed token sequence. + ModelOutput forward(const torch::Tensor& tokens, + const torch::Tensor& positions, + std::vector& kv_caches, + const ModelInputParams& input_params) { + return model_(tokens, positions, kv_caches, input_params); + } + + // Compute logits for selected token positions. + // hidden_states: [total_tokens, hidden_size] + // selected_idxes: optional token selector + // returns: [selected_tokens, vocab_size] + torch::Tensor logits(const torch::Tensor& hidden_states, + const torch::Tensor& selected_idxes) { + torch::Tensor h = hidden_states; + if (selected_idxes.defined()) { + h = h.index_select(/*dim=*/0, selected_idxes); + } + return head_(h); + } + + // Embedding / rerank path: return selected hidden states as-is. + torch::Tensor pooler(const torch::Tensor& hidden_states, + const torch::Tensor& selected_idxes) { + torch::Tensor h = hidden_states; + if (selected_idxes.defined()) { + h = h.index_select(/*dim=*/0, selected_idxes); + } + return h; + } + + // Load weights from checkpoint. + // + // Supported checkpoint layouts: + // • BlinkDL native: flat keys (emb.weight, blocks.*.*, ln_out.*, + // head.weight) • HuggingFace: "model." prefix (model.emb.weight, …, + // model.ln_out.*, lm_head.weight) + void load_model(std::unique_ptr loader) { + for (const auto& state_dict : loader->get_state_dicts()) { + // Try to strip common top-level prefixes for the body weights. + // An empty prefix ("") matches the native BlinkDL flat format. + model_->load_state_dict(state_dict->get_dict_with_prefix( + std::vector{"model.", ""})); + + // LM head weight: "head.weight" (native) or "lm_head.weight" (HF) + for (const auto& key : std::vector{ + "head.weight", "lm_head.weight", "model.head.weight"}) { + torch::Tensor head_w = state_dict->get_tensor(key); + if (head_w.defined()) { + head_->weight = + head_w.to(head_->weight.device()).to(head_->weight.scalar_type()); + break; + } + } + } + } + + // Required by CausalLMImpl (no-ops for non-MoE models) + void prepare_expert_weight(int32_t /*layer_id*/, + const std::vector& /*expert_ids*/) {} + void update_expert_weight(int32_t /*layer_id*/) {} + + private: + RWKV7Model model_{nullptr}; + torch::nn::Linear head_{nullptr}; +}; + +TORCH_MODULE(RWKV7ForCausalLM); + +// --------------------------------------------------------------------------- +// Model registration +// --------------------------------------------------------------------------- + +REGISTER_CAUSAL_MODEL(rwkv7, RWKV7ForCausalLM); + +// Model args — supports the RWKV-7 HuggingFace config.json format. +// +// Minimal config.json example: +// { +// "model_type": "rwkv7", +// "hidden_size": 768, +// "num_hidden_layers": 12, +// "vocab_size": 65536, +// "head_size": 64 +// } +REGISTER_MODEL_ARGS(rwkv7, [&] { + LOAD_ARG_OR(model_type, "model_type", "rwkv7"); + LOAD_ARG_OR(dtype, "torch_dtype", ""); + + LOAD_ARG_OR(vocab_size, "vocab_size", 65536); + LOAD_ARG_OR(hidden_size, "hidden_size", 768); + LOAD_ARG_OR(n_layers, "num_hidden_layers", 12); + + // RWKV-7 head_size_a — stored as "head_size" in the config. + LOAD_ARG_OR(head_dim, "head_size", 64); + + // Derive n_heads from hidden_size / head_size. + LOAD_ARG_OR_FUNC(n_heads, "num_attention_heads", [&] { + return args->hidden_size() / args->head_dim(); + }); + + // FFN intermediate size (4 × hidden_size by default). + LOAD_ARG_OR_FUNC(intermediate_size, "intermediate_size", [&] { + return 4LL * args->hidden_size(); + }); + + // LayerNorm eps used for ln0 / ln1 / ln2 / ln_out (default 1e-5). + LOAD_ARG_OR(layer_norm_eps, "layer_norm_eps", 1e-5f); + + // RWKV-7 World checkpoints are trained with ctx4096. + LOAD_ARG_OR(max_position_embeddings, "max_position_embeddings", 4096); + + LOAD_ARG_OR(bos_token_id, "bos_token_id", 0); + LOAD_ARG_OR(eos_token_id, "eos_token_id", 0); + + // RWKV-7 has no standard transformer attention; use a single dummy KV head + // so the standard paged KV cache is allocated as small as possible. + SET_ARG(n_kv_heads, static_cast(1)); + + // Linear-attention KV cache configuration. + // + // conv_cache per slot: [1, head_size * (2*n_heads + n_heads)] = [1, 3*H] + // [0:H] → att_x_prev (time-mix shift state) + // [H:2H] → ffn_x_prev (channel-mix shift state) + // [2H:3H] → unused (artefact of the KVCacheShape formula) + // + // ssm_cache per slot: [n_heads, head_size, head_size] ← W-matrix state + LOAD_ARG_OR_FUNC(linear_num_key_heads, "num_attention_heads", [&] { + return args->n_heads(); // int64_t matches linear_num_key_heads type + }); + LOAD_ARG_OR_FUNC(linear_key_head_dim, "head_size", [&] { + return static_cast(args->head_dim()); + }); + LOAD_ARG_OR_FUNC(linear_num_value_heads, "num_attention_heads", [&] { + return static_cast(args->n_heads()); + }); + LOAD_ARG_OR_FUNC(linear_value_head_dim, "head_size", [&] { + return static_cast(args->head_dim()); + }); + // conv_state_len = linear_conv_kernel_dim - 1 = 1 (one step of history) + SET_ARG(linear_conv_kernel_dim, static_cast(2)); + + // Set full_attention_interval = n_layers + 1 so that: + // + // 1. has_linear_attention_layers(args) == true + // (because full_attention_interval > 1) + // → worker_impl sets enable_linear_attention = true + // → KVCacheShape allocates conv_cache + ssm_cache tensors + // + // 2. is_linear_attention_layer(i, full_attention_interval) == true + // for ALL layer indices i ∈ [0, n_layers-1] + // (because (i+1) % (n_layers+1) ≠ 0 for i < n_layers) + // → create_kv_cache_impl creates LinearAttentionKVCacheImpl per layer + // + // Without this, KVCacheImpl (standard paged blocks) would be created for + // every layer, and get_conv_cache() / get_ssm_cache() would return undefined + // tensors, causing a DCHECK failure at the first forward pass. + SET_ARG(full_attention_interval, static_cast(args->n_layers() + 1)); + + // Mark all layers as "rwkv7" so that has_linear_attention_layers() also + // returns true via the layer_types path, providing extra clarity. + SET_ARG( + layer_types, + std::vector(static_cast(args->n_layers()), "rwkv7")); + + SET_ARG(stop_token_ids, std::unordered_set({args->eos_token_id()})); +}); + +// Tokenizer args. +// +// RWKV-7 World models use the official trie tokenizer backed by +// rwkv_vocab_v20230424.txt (exported by tools/convert_rwkv7_world.py). +REGISTER_TOKENIZER_ARGS(rwkv7, [&] { + SET_ARG(tokenizer_type, "rwkv"); + SET_ARG(vocab_file, "rwkv_vocab_v20230424.txt"); +}); + +} // namespace xllm diff --git a/xllm/models/models.h b/xllm/models/models.h index a933a54041..acf31ba40f 100644 --- a/xllm/models/models.h +++ b/xllm/models/models.h @@ -98,6 +98,7 @@ limitations under the License. #include "llm/qwen2.h" // IWYU pragma: keep #include "llm/qwen3.h" // IWYU pragma: keep #include "llm/qwen3_moe.h" // IWYU pragma: keep +#include "llm/rwkv7.h" // IWYU pragma: keep #include "vlm/qwen2_5_vl.h" // IWYU pragma: keep #include "vlm/qwen2_vl.h" // IWYU pragma: keep #include "vlm/qwen3_vl.h" // IWYU pragma: keep