Upload wkv.py

Trained_20G/wkv.py

@@ -1,3 +1,4 @@
+import os
 import torch
 from einops import rearrange
 
@@ -41,6 +42,9 @@ else:
     def compile_decorator(func):
         return func
 
+wkv_mode = os.environ.get("WKV_MODE", "fused")
+wkv_mode = wkv_mode.lower()
+assert wkv_mode in ['fused', 'chunk', 'pytorch']
 
 class Rwkv_Tmix_x070(nn.Module):
     def __init__(self, args: RwkvHybridConfig, layer_id, **kwargs):
@@ -205,23 +209,22 @@ class Rwkv_Tmix_x070(nn.Module):
         output_final_state,
         cu_seqlens
     ):
-        if r.device.type == "cpu":
+        if wkv_mode == 'pytorch':
             r, w, k, v, a, b = map(lambda x: rearrange(
                 x, 'b l (h d) -> b h l d', h=self.n_head), (r, w, k, v, a, b))
             o, state = native_recurrent_rwkv7(
                 r=r, k=k, v=v, w=w,
                 a=a, b=b,
                 scale=1.0,
-                initial_state=s.transpose(-1, -2),
+                initial_state=s,
                 output_final_state=True,
                 head_first=True,
             )
-            state = state.transpose(-1, -2)
             x = rearrange(o, "b h l d -> b l (h d)")
         else:
             r, w, k, v, a, b = map(lambda x: rearrange(
                 x, 'b l (h d) -> b l h d', h=self.n_head), (r, w, k, v, a, b))
-            wkv7_func = chunk_rwkv7 if r.shape[1] != 1 else fused_recurrent_rwkv7
+            wkv7_func = chunk_rwkv7 if wkv_mode == 'chunk' else fused_recurrent_rwkv7
             o, state = wkv7_func(
                 r=r, k=k, v=v, w=w,
                 a=a, b=b,