tvm icon indicating copy to clipboard operation
tvm copied to clipboard

[Relax][Runtime] RNNState for Space State Models

Open Hzfengsy opened this issue 1 year ago • 2 comments

This commit adds the RNNState class to the Relax VM, similar to the PagedKVCache, for space state models like RWKV and mamba

cc @tqchen @MasterJH5574

Hzfengsy avatar Feb 14 '24 04:02 Hzfengsy

Refactored. Now AttentionKVCache and RNNState have the same base class. would be great if @MasterJH5574 @tqchen Can take a look

Hzfengsy avatar Feb 19 '24 07:02 Hzfengsy

please fix the lint error

tqchen avatar Feb 19 '24 13:02 tqchen

CI is green now :)

Hzfengsy avatar Feb 21 '24 12:02 Hzfengsy