torch.compiled version, fp64 chunklen128

RWKV-v5/src/module/TimeMix.py

@@ -320,7 +320,7 @@ def _forward_nocuda_optimized(self, x, last_state: tuple[torch.Tensor,torch.Tens
         shift_state_out = x[:,-1]
 
         # 24 is optimal chunk length (longer will use too much memory and cause precision problems or even numerical instability, shorter is inefficient)
-        chunk_len = 24
+        chunk_len = 128
 
         # padding to support fast path for non-exact chunk size multiple sequence lengths        
         n_padding = (chunk_len - x.size(-2) % chunk_len) % chunk_len

RWKV-v5/src/module/rwkv_inner.py

@@ -5,7 +5,10 @@
 from torch import Tensor
 
 # 24 is optimal chunk length (longer will use too much memory and cause precision problems or even numerical instability, shorter is inefficient)
-def rwkv_inner(r,k,v,w,u,kv_state,chunk_len:int=24,precision_dtype:torch.dtype=torch.float32):
+@torch.jit.ignore
+@torch.compile
+def rwkv_inner(r,k,v,w,u,kv_state,chunk_len:int=24):
+    precision_dtype:torch.dtype=torch.float64
     """
     expects
     r : (B,H,L,K)
@@ -38,7 +41,7 @@ def rwkv_inner(r,k,v,w,u,kv_state,chunk_len:int=24,precision_dtype:torch.dtype=t
         if precision_dtype == torch.float32:
             precision_min_val = 0.005 # good for fp32 (1.175e-38 ^ (1/16.0) < 0.00426)
         else: #elif precision_dtype == torch.float64:
-            precision_min_val = 1e-10 # good for fp64 (1.7e-308 ^ (1/16.0) < 5.8e-20)
+            precision_min_val = 0.005 # to match fp32
         w = w.clamp(precision_min_val)
 
         # calculate cumulative decay in log space where it won't overflow