├── .gitignore ├── README.md ├── flash_rwkv ├── __init__.py ├── rwkv5 │ ├── __init__.py │ ├── wkv5_kernel.cu │ ├── wkv5_op.cpp │ └── wkv5_op.py └── rwkv6 │ ├── __init__.py │ ├── wkv6_kernel.cu │ ├── wkv6_op.cpp │ └── wkv6_op.py ├── profile ├── profile_rwkv6_linear_attention.py └── rwkv6_profile_result │ ├── chunk │ ├── worker0.1714447640486073830.pt.trace.json │ ├── worker0.1714447640592101360.pt.trace.json │ └── worker0.1714447640627745841.pt.trace.json │ ├── cuda │ ├── worker0.1714573540025583042.pt.trace.json │ ├── worker0.1714573540049675537.pt.trace.json │ └── worker0.1714573540070875800.pt.trace.json │ ├── hf │ ├── worker0.1714447608047920646.pt.trace.json │ ├── worker0.1714447608389966965.pt.trace.json │ └── worker0.1714447608630063783.pt.trace.json │ └── recurrent │ ├── worker0.1714447661248982960.pt.trace.json │ ├── worker0.1714447661534350082.pt.trace.json │ └── worker0.1714447661631554849.pt.trace.json ├── setup.py └── tests ├── test_base_linear_attention.py ├── test_rwkv5_linear_attention.py └── test_rwkv6_linear_attention.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/.gitignore -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/README.md -------------------------------------------------------------------------------- /flash_rwkv/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/flash_rwkv/__init__.py -------------------------------------------------------------------------------- /flash_rwkv/rwkv5/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/flash_rwkv/rwkv5/__init__.py -------------------------------------------------------------------------------- /flash_rwkv/rwkv5/wkv5_kernel.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/flash_rwkv/rwkv5/wkv5_kernel.cu -------------------------------------------------------------------------------- /flash_rwkv/rwkv5/wkv5_op.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/flash_rwkv/rwkv5/wkv5_op.cpp -------------------------------------------------------------------------------- /flash_rwkv/rwkv5/wkv5_op.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/flash_rwkv/rwkv5/wkv5_op.py -------------------------------------------------------------------------------- /flash_rwkv/rwkv6/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/flash_rwkv/rwkv6/__init__.py -------------------------------------------------------------------------------- /flash_rwkv/rwkv6/wkv6_kernel.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/flash_rwkv/rwkv6/wkv6_kernel.cu -------------------------------------------------------------------------------- /flash_rwkv/rwkv6/wkv6_op.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/flash_rwkv/rwkv6/wkv6_op.cpp -------------------------------------------------------------------------------- /flash_rwkv/rwkv6/wkv6_op.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/flash_rwkv/rwkv6/wkv6_op.py -------------------------------------------------------------------------------- /profile/profile_rwkv6_linear_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/profile_rwkv6_linear_attention.py -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/chunk/worker0.1714447640486073830.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/chunk/worker0.1714447640486073830.pt.trace.json -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/chunk/worker0.1714447640592101360.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/chunk/worker0.1714447640592101360.pt.trace.json -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/chunk/worker0.1714447640627745841.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/chunk/worker0.1714447640627745841.pt.trace.json -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/cuda/worker0.1714573540025583042.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/cuda/worker0.1714573540025583042.pt.trace.json -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/cuda/worker0.1714573540049675537.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/cuda/worker0.1714573540049675537.pt.trace.json -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/cuda/worker0.1714573540070875800.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/cuda/worker0.1714573540070875800.pt.trace.json -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/hf/worker0.1714447608047920646.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/hf/worker0.1714447608047920646.pt.trace.json -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/hf/worker0.1714447608389966965.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/hf/worker0.1714447608389966965.pt.trace.json -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/hf/worker0.1714447608630063783.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/hf/worker0.1714447608630063783.pt.trace.json -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/recurrent/worker0.1714447661248982960.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/recurrent/worker0.1714447661248982960.pt.trace.json -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/recurrent/worker0.1714447661534350082.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/recurrent/worker0.1714447661534350082.pt.trace.json -------------------------------------------------------------------------------- /profile/rwkv6_profile_result/recurrent/worker0.1714447661631554849.pt.trace.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/profile/rwkv6_profile_result/recurrent/worker0.1714447661631554849.pt.trace.json -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/setup.py -------------------------------------------------------------------------------- /tests/test_base_linear_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/tests/test_base_linear_attention.py -------------------------------------------------------------------------------- /tests/test_rwkv5_linear_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/tests/test_rwkv5_linear_attention.py -------------------------------------------------------------------------------- /tests/test_rwkv6_linear_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BBuf/flash-rwkv/HEAD/tests/test_rwkv6_linear_attention.py --------------------------------------------------------------------------------