Upload 11 files

config.json

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

configuration_rwkv5.py

@@ -101,6 +101,7 @@ 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,
@@ -114,6 +115,7 @@ 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

cpp_kernels.py

@@ -0,0 +1,55 @@
+from torch.utils import cpp_extension
+import pathlib
+import os
+import subprocess
+
+def _get_cuda_bare_metal_version(cuda_dir):
+    raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"],
+                                         universal_newlines=True)
+    output = raw_output.split()
+    release_idx = output.index("release") + 1
+    release = output[release_idx].split(".")
+    bare_metal_major = release[0]
+    bare_metal_minor = release[1][0]
+
+    return raw_output, bare_metal_major, bare_metal_minor
+
+def _create_build_dir(buildpath):
+    try:
+        os.mkdir(buildpath)
+    except OSError:
+        if not os.path.isdir(buildpath):
+            print(f"Creation of the build directory {buildpath} failed")
+
+# Check if cuda 11 is installed for compute capability 8.0
+cc_flag = []
+_, bare_metal_major, bare_metal_minor = _get_cuda_bare_metal_version(cpp_extension.CUDA_HOME)
+if int(bare_metal_major) >= 11:
+    cc_flag.append('-gencode')
+    cc_flag.append('arch=compute_80,code=sm_80')
+    if int(bare_metal_minor) >= 7:
+        cc_flag.append('-gencode')
+        cc_flag.append('arch=compute_90,code=sm_90')
+
+# Build path
+srcpath = pathlib.Path(__file__).parent.absolute()
+buildpath = srcpath / 'build'
+_create_build_dir(buildpath)
+
+def _cpp_extention_load_helper(name, sources, extra_cuda_flags):
+    return cpp_extension.load(
+        name=name,
+        sources=sources,
+        build_directory=buildpath,
+        extra_cflags=['-O3', ],
+        extra_cuda_cflags=['-O3',
+                           '-gencode', 'arch=compute_70,code=sm_70',
+                           '--use_fast_math'] + extra_cuda_flags + cc_flag,
+        verbose=1
+    )
+
+extra_flags = []
+
+cache_wkv5_sources = ["./rwkv5_op.cpp",
+           "./rwkv5.cu"]
+cache_wkv5 = _cpp_extention_load_helper("cache_wkv5", cache_wkv5_sources, extra_flags)

modeling_rwkv5.py

@@ -36,6 +36,7 @@ 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"
@@ -45,42 +46,29 @@ RWKV_PRETRAINED_MODEL_ARCHIVE_LIST = [
     
 ]
 
-rwkv5_cuda_kernel = None
-
-def load_wkv5_cuda_kernel(config):
-    global rwkv5_cuda_kernel
-    if config.model_version == "5_2" and torch.cuda.is_available():
-        HEAD_SIZE = args.attention_hidden_size // args.head_size
-        module_root = pathlib.Path(__file__).parent
-        rwkv5_cuda_kernel = load(name="rwkv5", sources=[f"{module_root}/rwkv5_op.cpp", f"{module_root}/rwkv5.cu"],
-            verbose=True, extra_cuda_cflags=["-res-usage", "--use_fast_math", "-O3", "-Xptxas -O3" if os.name != "nt" else "", "--extra-device-vectorization", f"-D_N_={HEAD_SIZE}"])
-
-class RWKV_5(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, B, T, C, H, state, r, k, v, w, u):
-        with torch.no_grad():
-            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:
-                rwkv5_cuda_kernel.forward_bf16(B, T, C, H, state, r, k, v, w, u, y)
-            elif r.dtype == torch.float16:
-                rwkv5_cuda_kernel.forward_fp16(B, T, C, H, state, r, k, v, w, u, y)
-            elif r.dtype == torch.float32:
-                rwkv5_cuda_kernel.forward_fp32(B, T, C, H, state, r, k, v, w, u, y)
-            return y, state
+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)
@@ -153,12 +141,20 @@ class RwkvSelfAttention(nn.Module):
         super().__init__()
         self.config = config
         self.layer_id = layer_id
-        kernel_loaded = rwkv5_cuda_kernel is not None
-        if torch.cuda.is_available() and not kernel_loaded:
-            try:
-                load_wkv5_cuda_kernel(config)
-            except Exception:
-                logger.info("Could not load the custom CUDA kernel for RWKV5 attention.")
+        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
         # https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4neo/src/model.py#L146
         num_attention_heads = self.hidden_size // config.head_size
@@ -206,7 +202,7 @@ class RwkvSelfAttention(nn.Module):
             gate = hidden* self.time_mix_gate + shifted * (1 - self.time_mix_gate)
             gate = F.silu(self.gate(gate))
 
-        if rwkv5_cuda_kernel is None:
+        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)
@@ -235,8 +231,8 @@ 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 rwkv5_cuda_kernel is not None and seq_mode:
-                rwkv, layer_state = RWKV_5.apply(1, T, self.hidden_size, H, layer_state.transpose(-1, -2).contiguous(), 
+            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)