Upload 9 files

config.json

@@ -21,6 +21,5 @@
   "tie_word_embeddings": false,
   "transformers_version": "4.33.1",
   "use_cache": true,
-  "use_cache_kernel": true,
   "vocab_size": 65536
 }

configuration_rwkv5.py

@@ -101,7 +101,6 @@ class Rwkv5Config(PretrainedConfig):
         eos_token_id=0,
         rescale_every=6,
         tie_word_embeddings=False,
-        use_cache_kernel=True,
         use_cache=True,
         model_version="5_2",
         **kwargs,
@@ -115,7 +114,6 @@ class Rwkv5Config(PretrainedConfig):
         self.intermediate_size = None
         self.layer_norm_epsilon = layer_norm_epsilon
         self.rescale_every = rescale_every
-        self.use_cache_kernel = use_cache_kernel
         self.use_cache = use_cache
 
         self.bos_token_id = bos_token_id

generation_config.json

@@ -0,0 +1,12 @@
+{
+    "chat_format": "chatml",
+    "eos_token_id": 0,
+    "pad_token_id": 0,
+    "max_window_size": 4096,
+    "max_new_tokens": 4096,
+    "do_sample": true,
+    "top_k": 0,
+    "top_p": 0.1,
+    "repetition_penalty": 1.0,
+    "transformers_version": "4.31.1"
+}

modeling_rwkv5.py

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """PyTorch RWKV5 World model."""
+
 import math
 from dataclasses import dataclass
 from pathlib import Path
@@ -36,7 +37,8 @@ from transformers.utils import (
     logging,
 )
 from .configuration_rwkv5 import Rwkv5Config
-from .cpp_kernels import cache_wkv5
+
+
 logger = logging.get_logger(__name__)
 
 _CHECKPOINT_FOR_DOC = "RWKV/rwkv-5-world"
@@ -46,30 +48,6 @@ RWKV_PRETRAINED_MODEL_ARCHIVE_LIST = [
     
 ]
 
-def rwkv_linear_attention_v5_2_cuda(B, T, C, H, state, r, k, v, w, u, cache_kernels):
-    assert HEAD_SIZE == C // H
-    ctx.B = B
-    ctx.T = T
-    ctx.C = C
-    ctx.H = H
-    assert state.dtype == torch.float32
-    assert w.dtype == torch.float32
-    assert r.is_contiguous()
-    assert k.is_contiguous()
-    assert v.is_contiguous()
-    assert w.is_contiguous()                            
-    assert u.is_contiguous()                            
-    assert state.is_contiguous()
-
-    y = torch.empty((B, T, C), device=w.device, dtype=r.dtype, memory_format=torch.contiguous_format)
-    if r.dtype == torch.bfloat16:
-        cache_kernels.forward_bf16(B, T, C, H, state, r, k, v, w, u, y)
-    elif r.dtype == torch.float16:
-        cache_kernels.forward_fp16(B, T, C, H, state, r, k, v, w, u, y)
-    elif r.dtype == torch.float32:
-        cache_kernels.forward_fp32(B, T, C, H, state, r, k, v, w, u, y)
-    return y, state
-
 def rwkv_linear_attention_v5_0(H, S, T, hidden, time_decay, time_first, receptance, key, value, lxw, lxb, ow, state, return_state=False, seq_mode=True):
     time_decay = torch.exp(-torch.exp(time_decay.float())).reshape(-1,1,1)
     time_first = torch.exp(time_first.float()).reshape(-1,1,1)
@@ -107,7 +85,7 @@ def rwkv_linear_attention_v5_0(H, S, T, hidden, time_decay, time_first, receptan
 
     return out, state
 
-def rwkv_linear_attention_v5_2_cpu(H, S, T, n_head, hidden, time_decay, time_first, receptance, key, value, gate, lxw, lxb, ow, state, return_state=False, seq_mode=True):
+def rwkv_linear_attention_v5_2(H, S, T, n_head, hidden, time_decay, time_first, receptance, key, value, gate, lxw, lxb, ow, state, return_state=False, seq_mode=True):
     time_decay = torch.exp(-torch.exp(time_decay.float())).reshape(-1,1,1).reshape(n_head, -1, 1)
     time_first = time_first.float().reshape(-1,1,1).reshape(n_head, -1, 1)
     lxw = lxw.float()
@@ -136,55 +114,43 @@ def rwkv_linear_attention_v5_2_cpu(H, S, T, n_head, hidden, time_decay, time_fir
         out = out @ ow
 
     return out, state
+
+
 class RwkvSelfAttention(nn.Module):
     def __init__(self, config, layer_id=0):
         super().__init__()
         self.config = config
         self.layer_id = layer_id
-        if config.use_cache_kernel:
-            # pre check if the support files existing
-            module_root = pathlib.Path(__file__).parent
-            src_files = ("rwkv5_op.cpp", "rwkv5.cu")
-            if any(not (module_root/src).is_file() for src in src_files):
-                warnings.warn("State cache kernel source files (.cpp and .cu) not found.")
-                self.cache_kernels = None
-            else:
-                try:
-                    from .cpp_kernels import cache_wkv5
-                    self.cache_kernels = cache_wkv5
-                except ImportError:
-                    warnings.warn("Failed to import KV cache kernels.")
-                    self.cache_kernels = None
-        self.hidden_size = config.hidden_size
+        hidden_size = config.hidden_size
         # https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4neo/src/model.py#L146
-        num_attention_heads = self.hidden_size // config.head_size
+        num_attention_heads = hidden_size // config.head_size
         self.num_attention_heads = num_attention_heads
         attention_hidden_size = (
-            config.attention_hidden_size if config.attention_hidden_size is not None else self.hidden_size
+            config.attention_hidden_size if config.attention_hidden_size is not None else hidden_size
         )
         self.attention_hidden_size = attention_hidden_size
 
         if self.config.model_version == "5_2":
             self.time_decay = nn.Parameter(torch.empty(num_attention_heads, config.head_size))
             self.time_faaaa = nn.Parameter(torch.empty(num_attention_heads, config.head_size))
-            self.time_mix_gate = nn.Parameter(torch.empty(1, 1, self.hidden_size))
+            self.time_mix_gate = nn.Parameter(torch.empty(1, 1, hidden_size))
         else:
             self.time_decay = nn.Parameter(torch.empty(num_attention_heads))
             self.time_first = nn.Parameter(torch.empty(num_attention_heads))
 
-        self.time_mix_key = nn.Parameter(torch.empty(1, 1, self.hidden_size))
-        self.time_mix_value = nn.Parameter(torch.empty(1, 1, self.hidden_size))
-        self.time_mix_receptance = nn.Parameter(torch.empty(1, 1, self.hidden_size))
+        self.time_mix_key = nn.Parameter(torch.empty(1, 1, hidden_size))
+        self.time_mix_value = nn.Parameter(torch.empty(1, 1, hidden_size))
+        self.time_mix_receptance = nn.Parameter(torch.empty(1, 1, hidden_size))
 
         self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
-        self.key = nn.Linear(self.hidden_size, attention_hidden_size, bias=False)
-        self.value = nn.Linear(self.hidden_size, attention_hidden_size, bias=False)
-        self.receptance = nn.Linear(self.hidden_size, attention_hidden_size, bias=False)
+        self.key = nn.Linear(hidden_size, attention_hidden_size, bias=False)
+        self.value = nn.Linear(hidden_size, attention_hidden_size, bias=False)
+        self.receptance = nn.Linear(hidden_size, attention_hidden_size, bias=False)
         if self.config.model_version == "5_2":
-            self.gate = nn.Linear(self.hidden_size, attention_hidden_size, bias=False)
-        self.output = nn.Linear(attention_hidden_size, self.hidden_size, bias=False)
+            self.gate = nn.Linear(hidden_size, attention_hidden_size, bias=False)
+        self.output = nn.Linear(attention_hidden_size, hidden_size, bias=False)
         # https://github.com/BlinkDL/RWKV-LM/blob/3db37a72356b736966ddd377268f02b80963af3f/RWKV-v4neo/src/model.py#L190C1-L190C1
-        self.ln_x = nn.GroupNorm(self.hidden_size // config.head_size, self.hidden_size)
+        self.ln_x = nn.GroupNorm(hidden_size // config.head_size, hidden_size)
 
     # TODO: maybe jit, otherwise move inside forward
     def extract_key_value(self, H, S, T, hidden, state=None):
@@ -200,18 +166,19 @@ class RwkvSelfAttention(nn.Module):
         receptance = hidden * self.time_mix_receptance + shifted * (1 - self.time_mix_receptance)
         if self.config.model_version == "5_2":
             gate = hidden* self.time_mix_gate + shifted * (1 - self.time_mix_gate)
-            gate = F.silu(self.gate(gate))
 
-        if self.cache_kernels is None:
-            if hidden.size(1) == 1 and state is not None:
-                receptance = self.receptance(receptance).to(torch.float32).view(H, 1, S)
-                key = self.key(key).to(torch.float32).view(H, S, 1)
-                value = self.value(value).to(torch.float32).view(H, 1, S)
-            else:
-                # https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L693
-                key = self.key(key).to(torch.float32).view(T, H, S).transpose(0, 1).transpose(-2, -1)
-                value = self.value(value).to(torch.float32).view(T, H, S).transpose(0, 1)
-                receptance = self.receptance(receptance).to(torch.float32).view(T, H, S).transpose(0, 1)
+        if hidden.size(1) == 1 and state is not None:
+            receptance = self.receptance(receptance).to(torch.float32).view(H, 1, S)
+            key = self.key(key).to(torch.float32).view(H, S, 1)
+            value = self.value(value).to(torch.float32).view(H, 1, S)
+        else:
+            # https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L693
+            key = self.key(key).to(torch.float32).view(T, H, S).transpose(0, 1).transpose(-2, -1)
+            value = self.value(value).to(torch.float32).view(T, H, S).transpose(0, 1)
+            receptance = self.receptance(receptance).to(torch.float32).view(T, H, S).transpose(0, 1)
+
+        if self.config.model_version == "5_2":
+            gate = F.silu(self.gate(gate))
 
         if state is not None:
             state[0][:, :, self.layer_id] = hidden[:, -1]
@@ -231,34 +198,25 @@ class RwkvSelfAttention(nn.Module):
             receptance, key, value, state = self.extract_key_value(H, S, T, hidden, state=state)
         layer_state = state[1][:, :, :, :, self.layer_id] if state is not None else None
         if self.config.model_version == "5_2":
-            if self.cache_kernels is not None and seq_mode:
-                rwkv, layer_state = rwkv_linear_attention_v5_2_cuda(1, T, self.hidden_size, H, layer_state.transpose(-1, -2).contiguous(), 
-                    receptance, key, value, self.time_decay, self.time_faaaa,)
-                layer_state = layer_state.transpose(-1,-2)
-                rwkv = rwkv.reshape(T, H*N)
-                rwkv = F.group_norm(rwkv, num_groups=H, weight=self.ln_x.weight, bias=self.ln_x.bias)
-                rwkv = rwkv.to(dtype=hidden.dtype) * gate
-                rwkv = rwkv @ self.output.weight.t()
-            else:
-                rwkv, layer_state = rwkv_linear_attention_v5_2_cpu(
-                    H,
-                    S,
-                    T,
-                    self.num_attention_heads,
-                    hidden,
-                    self.time_decay,
-                    self.time_faaaa,
-                    receptance,
-                    key,
-                    value,
-                    gate,
-                    self.ln_x.weight,
-                    self.ln_x.bias,
-                    self.output.weight.t(),
-                    state=layer_state,
-                    return_state=use_cache,
-                    seq_mode=seq_mode,
-                )
+            rwkv, layer_state = rwkv_linear_attention_v5_2(
+            H,
+            S,
+            T,
+            self.num_attention_heads,
+            hidden,
+            self.time_decay,
+            self.time_faaaa,
+            receptance,
+            key,
+            value,
+            gate,
+            self.ln_x.weight,
+            self.ln_x.bias,
+            self.output.weight.t(),
+            state=layer_state,
+            return_state=use_cache,
+            seq_mode=seq_mode,
+        )
         else:
             rwkv, layer_state = rwkv_linear_attention_v5_0(
                 H,

tokenization_rwkv_world.py

@@ -52,186 +52,52 @@ if TYPE_CHECKING:
 logger = logging.get_logger(__name__)
 
 VOCAB_FILES_NAMES = {
-    "vocab_file": "rwkv_vocab_v20230424.json",
+    "vocab_file": "rwkv_vocab_v20230424.txt",
 }
 
-
-class DATrie:
-    class Node:
-        def __init__(self, is_leaf=False, leaf_data=None, tail=""):
-            self._is_leaf = is_leaf
-            self._leaf_data = leaf_data
-            self._tail = tail
-            self._next_map = {}
-
-        def is_leaf(self):
-            return self._is_leaf
-
-        def set_leaf(self):
-            self._is_leaf = True
-
-        def has_next(self, w):
-            if w in self._next_map:
-                return True
-            return False
-
-        def add_node(self, w, node):
-            self._next_map[w] = node
-
-        def get_node(self, w):
-            if w in self._next_map:
-                return self._next_map[w]
-            return None
-
-        def get_tail(self):
-            return self._tail
-
-        def get_data(self):
-            return self._leaf_data
-
-        def set_data(self, data):
-            self._leaf_data = data
-
-    def __init__(self, special_ids):
-        self.root = self.Node()
-        self.data = {}
-        self.r_data = {}
-        self.special_ids = special_ids
-
-    def insert(self, word, data):
-        self.data[word] = data
-        self.r_data[data] = word
-        idx = 0
-        node = self.root
-        while idx < len(word):
-            w = word[idx]
-            is_leaf = (idx == (len(word) - 1))
-            leaf_data = (data if is_leaf else None)
-            # 不存在则插入
-            if not node.has_next(w):
-                node.add_node(w, self.Node(is_leaf=is_leaf, leaf_data=leaf_data))
-                # last word
-            node = node.get_node(w)
-            idx += 1
-        if not node.is_leaf():
-            node.set_leaf()
-            node.set_data(data)
-
-    def findStrict(self, word):
-        idx = 0
-        node = self.root
-        while node is not None and idx < len(word):
-            w = word[idx]
-            if not node.has_next(w):
-                return None
-                # last word
-            node = node.get_node(w)
-            idx += 1
-        if node.is_leaf():
-            return node.get_data()
-        return None
-
-    def prefix(self, word):
-        idx = 0
-        node = self.root
-        result = []
-        while node is not None and idx < len(word):
-            w = word[idx]
-            if not node.has_next(w):
-                return result
-                # last word
-            node = node.get_node(w)
-            if node.is_leaf():
-                result.append([word[:idx + 1], node.get_data()])
-            idx += 1
-        return result
-
-    def max_prefix(self, content, start_idx):
-        idx = start_idx
-        node = self.root
-        l = len(content)
-        result = [["", ], ]
-        while node is not None and idx < l:
-            w = content[idx]
-            if not node.has_next(w):
-                return result[-1]
-                # last word
-            node = node.get_node(w)
-            if node.is_leaf():
-                result.append([content[start_idx:idx + 1], node.get_data()])
+class TRIE:
+    __slots__ = tuple("ch,to,values,front".split(","))
+    to:list
+    values:set
+    def __init__(self, front=None, ch=None):
+        self.ch = ch
+        self.to = [None for ch in range(256)]
+        self.values = set()
+        self.front = front
+
+    def __repr__(self):
+        fr = self
+        ret = []
+        while(fr!=None):
+            if(fr.ch!=None):
+                ret.append(fr.ch)
+            fr = fr.front
+        return "<TRIE %s %s>"%(ret[::-1], self.values)
+    
+    def add(self, key:bytes, idx:int=0, val=None):
+        if(idx == len(key)):
+            if(val is None):
+                val = key
+            self.values.add(val)
+            return self
+        ch = key[idx]
+        if(self.to[ch] is None):
+            self.to[ch] = TRIE(front=self, ch=ch)
+        return self.to[ch].add(key, idx=idx+1, val=val)
+    
+    def find_longest(self, key:bytes, idx:int=0):
+        u:TRIE = self
+        ch:int = key[idx]
+        
+        while(u.to[ch] is not None):
+            u = u.to[ch]
             idx += 1
-        return result[-1]
-
-    def max_score(self, content, start_idx):
-        idx = start_idx
-        node = self.root
-        l = len(content)
-        result = [["", (3, 0)], ]
-        while node is not None and idx < l:
-            w = content[idx]
-            if not node.has_next(w):
-                break
-                # last word
-            node = node.get_node(w)
-            if node.is_leaf():
-                result.append([content[start_idx:idx + 1], node.get_data()])
-            idx += 1
-        if len(result) > 1:
-            result = sorted(result, key=lambda x: x[1][1])
-        return result[-1]
-
-    def match(self, content, add_unk=True, unk_id=-1, **kwargs):
-        # length
-        l = len(content)
-        i = 0
-        result_list = []
-        while i < l:
-            match_word = self.max_prefix(content=content, start_idx=i)
-            # print(match_word)
-            w = match_word[0]
-            if len(w) > 0:
-                result_list.append(match_word[1])
-                i += len(w)
-            else:
-                if add_unk:
-                    result_list.append(unk_id)
-                i += 1
-        return result_list
-
-    def id2str(self, ids, escape_special_ids=True, end_ids=[], **kwargs):
-        res_str = ""
-        for rid in ids:
-            if rid in self.r_data:
-                if rid in end_ids:
-                    break
-                if escape_special_ids and rid in self.special_ids:
-                    continue
-                rstr = self.r_data[rid]
-                res_str += rstr
-            elif rid == 0:
+            if(u.values):
+                ret = idx, u, u.values
+            if(idx==len(key)):
                 break
-            else:
-                print("ERROR unknown id %d" % rid)
-                res_str += "UNK"
-        return res_str
-
-    def id2str_v2(self, ids, escape_special_ids=True, end_ids=[], **kwargs):
-        res_str = ""
-        for rid in ids:
-            if rid in self.r_data:
-                if rid in end_ids:
-                    break
-                rstr = self.r_data[rid]
-                if escape_special_ids and rid in self.special_ids:
-                    continue
-                res_str += rstr
-            elif rid == 0:
-                break
-            else:
-                print("ERROR unknown id %d" % rid)
-                res_str += "UNK"
-        return res_str
-
+            ch = key[idx]
+        return ret
 
 class RWKVWorldTokenizer(PreTrainedTokenizer):
     vocab_files_names = VOCAB_FILES_NAMES
@@ -244,17 +110,30 @@ class RWKVWorldTokenizer(PreTrainedTokenizer):
             **kwargs
     ):
         self.add_bos_token = False
+        self.encoder = {}
+        sorted = [] # must be already sorted
+        with open(vocab_file, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+        for l in lines:
+            idx = int(l[:l.index(' ')])
+            x = eval(l[l.index(' '):l.rindex(' ')])
+            x = x.encode("utf-8") if isinstance(x, str) else x
+            assert isinstance(x, bytes)
+            assert len(x) == int(l[l.rindex(' '):])
+            sorted += [x]
+            self.encoder[idx] = x
 
-        with open(vocab_file, encoding="utf-8") as vocab_handle:
-            self.encoder = json.load(vocab_handle)
         super().__init__(
             errors=errors,
             **kwargs,
         )
-        self.decoder = {v: k for k, v in self.encoder.items()}
-        self.trie = DATrie(self.all_special_ids)
-        for k, v in self.encoder.items():
-            self.trie.insert(k, v)
+        self.decoder = {}
+        for k,v in self.encoder.items():
+            self.decoder[v] = int(k)
+
+        self.trie = TRIE()
+        for t, i in self.decoder.items():
+            _ = self.trie.add(t, val=(t, i))
         self.errors = errors  # how to handle errors in decoding
         self.cache = {}
 
@@ -311,9 +190,23 @@ class RWKVWorldTokenizer(PreTrainedTokenizer):
             return [1] + ([0] * len(token_ids_0))
         return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1))
 
+    def encodeBytes(self, src:bytes):
+        idx:int = 0
+        tokens = []
+        while (idx < len(src)):
+            _idx:int = idx
+            idx, _, values = self.trie.find_longest(src, idx)
+            assert(idx != _idx)
+            _, token = next(iter(values))            
+            tokens.append(token)
+        return tokens
+    
+    def decodeBytes(self, tokens):
+        return b''.join(map(lambda i: self.encoder[i], tokens))
+
     def _tokenize(self, text, **kwargs):
         """Tokenize a string."""
-        return self.trie.match(text, unk_id=self.unk_token_id, **kwargs)
+        return self.encodeBytes(text.encode("utf-8"))
 
     def _decode(self,
                token_ids: Union[int, List[int], "np.ndarray", "torch.Tensor", "tf.Tensor"],
@@ -326,13 +219,9 @@ class RWKVWorldTokenizer(PreTrainedTokenizer):
         if isinstance(token_ids, int):
             if token_ids in self.all_special_ids and skip_special_tokens:
                 return ""
-            return self.decoder.get(token_ids, self.unk_token)
+            return self.encoder.get(token_ids, self.unk_token)
         elif isinstance(token_ids, list):
-            return self.trie.id2str(
-                token_ids,
-                escape_special_ids=skip_special_tokens,
-                **kwargs
-            )
+            return self.decodeBytes(tokens).decode('utf-8')
         else:
             return token_ids
 
@@ -383,10 +272,10 @@ class RWKVWorldTokenizer(PreTrainedTokenizer):
     ) -> BatchEncoding:
         def get_input_ids(text):
             if isinstance(text, str):
-                text_id = self.trie.match(text, unk_id=self.unk_token_id)
+                text_id = self._tokenize(text)
                 return text_id
             elif isinstance(text, list) and len(text) > 0 and isinstance(text[0], str):
-                return [self.trie.match(t, unk_id=self.unk_token_id) for t in text]
+                return [self._tokenize(t) for t in text]
             elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], int):
                 return text
             else:
@@ -448,10 +337,10 @@ class RWKVWorldTokenizer(PreTrainedTokenizer):
     ) -> BatchEncoding:
         def get_input_ids(text):
             if isinstance(text, str):
-                text_id = self.trie.match(text, unk_id=self.unk_token_id)
+                text_id = self._tokenize(text)
                 return text_id
             elif isinstance(text, list) and len(text) > 0 and isinstance(text[0], str):
-                return [self.trie.match(t, unk_id=self.unk_token_id) for t in text]
+                return [self._tokenize(t) for t in text]
             elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], int):
                 return text
             else: