RWKV-8 ROSA-QKV-1bit Demo
This is using naive algorithm (not suffix automaton). Enter or randomize q/k/v (0/1 only), then click [Start Demo].
4 20
0.1 10
1def rosa_qkv_naive(qqq, kkk, vvv):
2 n=len(qqq); out=[-1]*n
3 for i in range(n):
4 for w in range(i+1,0,-1):
5 t=qqq[i+1-w:i+1]
6 for j in range(i-w,-1,-1):
7 if kkk[j:j+w]==t:
8 out[i]=vvv[j+w]
9 break
10 if out[i]!=-1:
11 break
12 return out
Fast version (click to expand)
1def rosa_qkv_ref_minus1(qqq, kkk, vvv): # note: input will never contain "-1"
2 n=len(qqq); y=[-1]*n; s=2*n+1; t=[None]*s; f=[-1]*s; m=[0]*s; r=[-1]*s; t[0]={}; g=0; u=1; w=h=0; assert n==len(kkk)==len(vvv)
3 for i,(q,k) in enumerate(zip(qqq,kkk)):
4 p,x=w,h
5 while p!=-1 and q not in t[p]: x=m[p] if x>m[p] else x; p=f[p]
6 p,x=(t[p][q],x+1) if p!=-1 else (0,0); v=p
7 while f[v]!=-1 and m[f[v]]>=x: v=f[v]
8 while v!=-1 and (m[v]<=0 or r[v]<0): v=f[v]
9 y[i]=vvv[r[v]+1] if v!=-1 else -1; w,h=p,x; j=u; u+=1; t[j]={}; m[j]=m[g]+1; p=g
10 while p!=-1 and k not in t[p]: t[p][k]=j; p=f[p]
11 if p==-1: f[j]=0
12 else:
13 d=t[p][k]
14 if m[p]+1==m[d]: f[j]=d
15 else:
16 b=u; u+=1; t[b]=t[d].copy(); m[b]=m[p]+1; f[b]=f[d]; r[b]=r[d]; f[d]=f[j]=b
17 while p!=-1 and t[p][k]==d: t[p][k]=b; p=f[p]
18 v=g=j
19 while v!=-1 and r[v]<i: r[v]=i; v=f[v]
20 return y