Kernels:
Trusted publisher
drbh commited on
Commit ·
2f27b37
unverified ·
0
Parent(s):
Migrated from kernels-community/rwkv
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .gitattributes +110 -0
- README.md +16 -0
- benchmarks/benchmark.py +81 -0
- build.toml +31 -0
- build/torch210-cu128-x86_64-windows/__init__.py +170 -0
- build/torch210-cu128-x86_64-windows/_ops.py +9 -0
- build/torch210-cu128-x86_64-windows/_rwkv_cuda_38ccc47.pyd +3 -0
- build/torch210-cu128-x86_64-windows/metadata.json +15 -0
- build/torch210-cu128-x86_64-windows/rwkv/__init__.py +26 -0
- build/torch210-cxx11-cu126-aarch64-linux/__init__.py +170 -0
- build/torch210-cxx11-cu126-aarch64-linux/_ops.py +9 -0
- build/torch210-cxx11-cu126-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so +3 -0
- build/torch210-cxx11-cu126-aarch64-linux/metadata.json +13 -0
- build/torch210-cxx11-cu126-aarch64-linux/rwkv/__init__.py +26 -0
- build/torch210-cxx11-cu126-x86_64-linux/__init__.py +170 -0
- build/torch210-cxx11-cu126-x86_64-linux/_ops.py +9 -0
- build/torch210-cxx11-cu126-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so +3 -0
- build/torch210-cxx11-cu126-x86_64-linux/metadata.json +13 -0
- build/torch210-cxx11-cu126-x86_64-linux/rwkv/__init__.py +26 -0
- build/torch210-cxx11-cu128-aarch64-linux/__init__.py +170 -0
- build/torch210-cxx11-cu128-aarch64-linux/_ops.py +9 -0
- build/torch210-cxx11-cu128-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so +3 -0
- build/torch210-cxx11-cu128-aarch64-linux/metadata.json +15 -0
- build/torch210-cxx11-cu128-aarch64-linux/rwkv/__init__.py +26 -0
- build/torch210-cxx11-cu128-x86_64-linux/__init__.py +170 -0
- build/torch210-cxx11-cu128-x86_64-linux/_ops.py +9 -0
- build/torch210-cxx11-cu128-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so +3 -0
- build/torch210-cxx11-cu128-x86_64-linux/metadata.json +15 -0
- build/torch210-cxx11-cu128-x86_64-linux/rwkv/__init__.py +26 -0
- build/torch210-cxx11-cu130-aarch64-linux/__init__.py +170 -0
- build/torch210-cxx11-cu130-aarch64-linux/_ops.py +9 -0
- build/torch210-cxx11-cu130-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so +3 -0
- build/torch210-cxx11-cu130-aarch64-linux/metadata.json +15 -0
- build/torch210-cxx11-cu130-aarch64-linux/rwkv/__init__.py +26 -0
- build/torch210-cxx11-cu130-x86_64-linux/__init__.py +170 -0
- build/torch210-cxx11-cu130-x86_64-linux/_ops.py +9 -0
- build/torch210-cxx11-cu130-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so +3 -0
- build/torch210-cxx11-cu130-x86_64-linux/metadata.json +15 -0
- build/torch210-cxx11-cu130-x86_64-linux/rwkv/__init__.py +26 -0
- build/torch211-cxx11-cu126-aarch64-linux/__init__.py +170 -0
- build/torch211-cxx11-cu126-aarch64-linux/_ops.py +9 -0
- build/torch211-cxx11-cu126-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so +3 -0
- build/torch211-cxx11-cu126-aarch64-linux/metadata.json +13 -0
- build/torch211-cxx11-cu126-aarch64-linux/rwkv/__init__.py +26 -0
- build/torch211-cxx11-cu126-x86_64-linux/__init__.py +170 -0
- build/torch211-cxx11-cu126-x86_64-linux/_ops.py +9 -0
- build/torch211-cxx11-cu126-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so +3 -0
- build/torch211-cxx11-cu126-x86_64-linux/metadata.json +13 -0
- build/torch211-cxx11-cu126-x86_64-linux/rwkv/__init__.py +26 -0
- build/torch211-cxx11-cu128-aarch64-linux/__init__.py +170 -0
.gitattributes
ADDED
|
@@ -0,0 +1,110 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
+
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
+
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
+
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
+
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
+
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
+
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
+
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
+
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
+
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
+
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
build/torch27-cxx11-cu118-x86_64-linux/rwkv/_rwkv_eb0e3e5_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
build/torch27-cxx11-cu126-x86_64-linux/rwkv/_rwkv_eb0e3e5_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 38 |
+
build/torch27-cxx11-cu128-x86_64-linux/rwkv/_rwkv_eb0e3e5_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 39 |
+
build/torch28-cxx11-cu126-x86_64-linux/rwkv/_rwkv_eb0e3e5_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 40 |
+
build/torch28-cxx11-cu128-x86_64-linux/rwkv/_rwkv_eb0e3e5_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 41 |
+
build/torch28-cxx11-cu129-x86_64-linux/rwkv/_rwkv_eb0e3e5_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 42 |
+
build/torch29-cxx11-cu126-x86_64-linux/rwkv/_rwkv_eb0e3e5_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 43 |
+
build/torch29-cxx11-cu128-x86_64-linux/rwkv/_rwkv_eb0e3e5_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 44 |
+
build/torch29-cxx11-cu130-x86_64-linux/rwkv/_rwkv_eb0e3e5_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 45 |
+
build/torch27-cxx11-cu118-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 46 |
+
build/torch27-cxx11-cu126-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 47 |
+
build/torch27-cxx11-cu128-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 48 |
+
build/torch28-cxx11-cu126-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 49 |
+
build/torch28-cxx11-cu128-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 50 |
+
build/torch28-cxx11-cu129-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 51 |
+
build/torch29-cxx11-cu126-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 52 |
+
build/torch29-cxx11-cu128-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 53 |
+
build/torch29-cxx11-cu130-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 54 |
+
build/torch210-cxx11-cu126-x86_64-linux/_rwkv_d5e72cf.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 55 |
+
build/torch210-cxx11-cu128-x86_64-linux/_rwkv_d5e72cf.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 56 |
+
build/torch210-cxx11-cu130-x86_64-linux/_rwkv_d5e72cf.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 57 |
+
build/torch28-cxx11-cu126-x86_64-linux/_rwkv_d5e72cf.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 58 |
+
build/torch28-cxx11-cu128-x86_64-linux/_rwkv_d5e72cf.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 59 |
+
build/torch28-cxx11-cu129-x86_64-linux/_rwkv_d5e72cf.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 60 |
+
build/torch29-cxx11-cu126-x86_64-linux/_rwkv_d5e72cf.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 61 |
+
build/torch29-cxx11-cu128-x86_64-linux/_rwkv_d5e72cf.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 62 |
+
build/torch29-cxx11-cu130-x86_64-linux/_rwkv_d5e72cf.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 63 |
+
build/torch210-cxx11-cu126-x86_64-linux/_rwkv_1aa1cb1.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 64 |
+
build/torch210-cxx11-cu128-x86_64-linux/_rwkv_1aa1cb1.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 65 |
+
build/torch210-cxx11-cu130-x86_64-linux/_rwkv_1aa1cb1.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 66 |
+
build/torch28-cxx11-cu126-x86_64-linux/_rwkv_1aa1cb1.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 67 |
+
build/torch28-cxx11-cu128-x86_64-linux/_rwkv_1aa1cb1.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 68 |
+
build/torch28-cxx11-cu129-x86_64-linux/_rwkv_1aa1cb1.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 69 |
+
build/torch29-cxx11-cu126-x86_64-linux/_rwkv_1aa1cb1.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 70 |
+
build/torch29-cxx11-cu128-x86_64-linux/_rwkv_1aa1cb1.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 71 |
+
build/torch29-cxx11-cu130-x86_64-linux/_rwkv_1aa1cb1.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 72 |
+
build/torch210-cxx11-cu126-x86_64-linux/_rwkv_86a3859.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 73 |
+
build/torch210-cxx11-cu128-x86_64-linux/_rwkv_86a3859.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 74 |
+
build/torch210-cxx11-cu130-x86_64-linux/_rwkv_86a3859.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 75 |
+
build/torch28-cxx11-cu126-x86_64-linux/_rwkv_86a3859.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 76 |
+
build/torch28-cxx11-cu128-x86_64-linux/_rwkv_86a3859.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 77 |
+
build/torch28-cxx11-cu129-x86_64-linux/_rwkv_86a3859.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 78 |
+
build/torch29-cxx11-cu126-x86_64-linux/_rwkv_86a3859.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 79 |
+
build/torch29-cxx11-cu128-x86_64-linux/_rwkv_86a3859.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 80 |
+
build/torch29-cxx11-cu130-x86_64-linux/_rwkv_86a3859.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 81 |
+
build/torch210-cxx11-cu126-x86_64-linux/_rwkv_44f2fa4.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 82 |
+
build/torch210-cxx11-cu128-x86_64-linux/_rwkv_44f2fa4.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 83 |
+
build/torch210-cxx11-cu130-x86_64-linux/_rwkv_44f2fa4.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 84 |
+
build/torch28-cxx11-cu126-x86_64-linux/_rwkv_44f2fa4.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 85 |
+
build/torch28-cxx11-cu128-x86_64-linux/_rwkv_44f2fa4.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 86 |
+
build/torch28-cxx11-cu129-x86_64-linux/_rwkv_44f2fa4.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 87 |
+
build/torch29-cxx11-cu126-x86_64-linux/_rwkv_44f2fa4.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 88 |
+
build/torch29-cxx11-cu128-x86_64-linux/_rwkv_44f2fa4.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 89 |
+
build/torch29-cxx11-cu130-x86_64-linux/_rwkv_44f2fa4.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 90 |
+
build/torch210-cxx11-cu126-aarch64-linux/_rwkv_cuda_efd954c.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 91 |
+
build/torch210-cxx11-cu128-aarch64-linux/_rwkv_cuda_efd954c.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 92 |
+
build/torch210-cxx11-cu130-aarch64-linux/_rwkv_cuda_efd954c.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 93 |
+
build/torch29-cxx11-cu126-aarch64-linux/_rwkv_cuda_efd954c.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 94 |
+
build/torch29-cxx11-cu128-aarch64-linux/_rwkv_cuda_efd954c.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 95 |
+
build/torch29-cxx11-cu130-aarch64-linux/_rwkv_cuda_efd954c.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 96 |
+
build/torch210-cu128-x86_64-windows/_rwkv_cuda_38ccc47.pyd filter=lfs diff=lfs merge=lfs -text
|
| 97 |
+
build/torch210-cxx11-cu126-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 98 |
+
build/torch210-cxx11-cu128-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 99 |
+
build/torch210-cxx11-cu130-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 100 |
+
build/torch211-cxx11-cu126-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 101 |
+
build/torch211-cxx11-cu128-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 102 |
+
build/torch211-cxx11-cu130-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 103 |
+
build/torch29-cxx11-cu129-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 104 |
+
build/torch210-cxx11-cu126-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 105 |
+
build/torch210-cxx11-cu128-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 106 |
+
build/torch210-cxx11-cu130-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 107 |
+
build/torch211-cxx11-cu126-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 108 |
+
build/torch211-cxx11-cu128-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 109 |
+
build/torch211-cxx11-cu130-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
| 110 |
+
build/torch29-cxx11-cu129-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so filter=lfs diff=lfs merge=lfs -text
|
README.md
ADDED
|
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
tags:
|
| 3 |
+
- kernels
|
| 4 |
+
---
|
| 5 |
+
|
| 6 |
+
RWKV kernel for transformers
|
| 7 |
+
### Performance
|
| 8 |
+
|
| 9 |
+
<img class="dark:hidden border border-gray-200 dark:border-gray-700 rounded-lg" src="media/benches_light_animation.svg" />
|
| 10 |
+
<img class="hidden dark:block border border-gray-200 dark:border-gray-700 rounded-lg" src="media/benches_dark_animation.svg" />
|
| 11 |
+
|
| 12 |
+
<img class="dark:hidden border border-gray-200 dark:border-gray-700 rounded-lg" src="media/benches_light_latency.svg" />
|
| 13 |
+
<img class="hidden dark:block border border-gray-200 dark:border-gray-700 rounded-lg" src="media/benches_dark_latency.svg" />
|
| 14 |
+
|
| 15 |
+
<img class="dark:hidden border border-gray-200 dark:border-gray-700 rounded-lg" src="media/benches_light_throughput.svg" />
|
| 16 |
+
<img class="hidden dark:block border border-gray-200 dark:border-gray-700 rounded-lg" src="media/benches_dark_throughput.svg" />
|
benchmarks/benchmark.py
ADDED
|
@@ -0,0 +1,81 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
|
| 3 |
+
from kernels.benchmark import Benchmark
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
def rwkv_wkv_reference(
|
| 7 |
+
w: torch.Tensor, u: torch.Tensor, k: torch.Tensor, v: torch.Tensor
|
| 8 |
+
) -> torch.Tensor:
|
| 9 |
+
B, T, C = k.shape
|
| 10 |
+
device = k.device
|
| 11 |
+
dtype = k.dtype
|
| 12 |
+
|
| 13 |
+
y = torch.zeros(B, T, C, device=device, dtype=dtype)
|
| 14 |
+
|
| 15 |
+
# State: accumulated numerator, denominator, and max exponent
|
| 16 |
+
aa = torch.zeros(B, C, device=device, dtype=torch.float32)
|
| 17 |
+
bb = torch.zeros(B, C, device=device, dtype=torch.float32)
|
| 18 |
+
pp = torch.full((B, C), -1e38, device=device, dtype=torch.float32)
|
| 19 |
+
|
| 20 |
+
w = w.float()
|
| 21 |
+
u = u.float()
|
| 22 |
+
|
| 23 |
+
for t in range(T):
|
| 24 |
+
kt = k[:, t, :].float() # [B, C]
|
| 25 |
+
vt = v[:, t, :].float() # [B, C]
|
| 26 |
+
|
| 27 |
+
# Output computation
|
| 28 |
+
ww = u + kt
|
| 29 |
+
p = torch.maximum(pp, ww)
|
| 30 |
+
e1 = torch.exp(pp - p)
|
| 31 |
+
e2 = torch.exp(ww - p)
|
| 32 |
+
y[:, t, :] = ((e1 * aa + e2 * vt) / (e1 * bb + e2)).to(dtype)
|
| 33 |
+
|
| 34 |
+
# State update (note: w + pp, not pp - w)
|
| 35 |
+
ww = w + pp
|
| 36 |
+
p = torch.maximum(ww, kt)
|
| 37 |
+
e1 = torch.exp(ww - p)
|
| 38 |
+
e2 = torch.exp(kt - p)
|
| 39 |
+
aa = e1 * aa + e2 * vt
|
| 40 |
+
bb = e1 * bb + e2
|
| 41 |
+
pp = p
|
| 42 |
+
|
| 43 |
+
return y
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
class RwkvBenchmark(Benchmark):
|
| 47 |
+
seed: int = 42
|
| 48 |
+
|
| 49 |
+
def setup(self):
|
| 50 |
+
B, T, C = 2, 64, 256
|
| 51 |
+
|
| 52 |
+
self.w = torch.randn(
|
| 53 |
+
C, device=self.device, dtype=torch.float32
|
| 54 |
+
).abs() # Decay should be positive
|
| 55 |
+
self.u = torch.randn(C, device=self.device, dtype=torch.float32)
|
| 56 |
+
self.k = torch.randn(B, T, C, device=self.device, dtype=torch.float32) * 0.1
|
| 57 |
+
self.v = torch.randn(B, T, C, device=self.device, dtype=torch.float32) * 0.1
|
| 58 |
+
self.out = torch.zeros(B, T, C, device=self.device, dtype=torch.float32)
|
| 59 |
+
|
| 60 |
+
def benchmark_base(self):
|
| 61 |
+
self.out.zero_()
|
| 62 |
+
self.kernel.forward(self.w, self.u, self.k, self.v, self.out)
|
| 63 |
+
|
| 64 |
+
def verify_base(self) -> torch.Tensor:
|
| 65 |
+
return rwkv_wkv_reference(self.w, self.u, self.k, self.v)
|
| 66 |
+
|
| 67 |
+
def setup_large(self):
|
| 68 |
+
B, T, C = 8, 256, 512
|
| 69 |
+
|
| 70 |
+
self.w = torch.randn(C, device=self.device, dtype=torch.float32).abs()
|
| 71 |
+
self.u = torch.randn(C, device=self.device, dtype=torch.float32)
|
| 72 |
+
self.k = torch.randn(B, T, C, device=self.device, dtype=torch.float32) * 0.1
|
| 73 |
+
self.v = torch.randn(B, T, C, device=self.device, dtype=torch.float32) * 0.1
|
| 74 |
+
self.out = torch.zeros(B, T, C, device=self.device, dtype=torch.float32)
|
| 75 |
+
|
| 76 |
+
def benchmark_large(self):
|
| 77 |
+
self.out.zero_()
|
| 78 |
+
self.kernel.forward(self.w, self.u, self.k, self.v, self.out)
|
| 79 |
+
|
| 80 |
+
def verify_large(self) -> torch.Tensor:
|
| 81 |
+
return rwkv_wkv_reference(self.w, self.u, self.k, self.v)
|
build.toml
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[general]
|
| 2 |
+
name = "rwkv"
|
| 3 |
+
universal = false
|
| 4 |
+
|
| 5 |
+
[torch]
|
| 6 |
+
src = [
|
| 7 |
+
"torch-ext/torch_binding.cpp",
|
| 8 |
+
]
|
| 9 |
+
|
| 10 |
+
[kernel.rwkv]
|
| 11 |
+
depends = ["torch"]
|
| 12 |
+
backend = "cuda"
|
| 13 |
+
cuda-capabilities = [
|
| 14 |
+
"8.0",
|
| 15 |
+
"8.9",
|
| 16 |
+
"9.0",
|
| 17 |
+
"10.0",
|
| 18 |
+
"12.0",
|
| 19 |
+
]
|
| 20 |
+
include = ["."]
|
| 21 |
+
src = [
|
| 22 |
+
"rwkv/wkv_cuda.cu",
|
| 23 |
+
"rwkv/wkv_cuda_bf16.cu",
|
| 24 |
+
]
|
| 25 |
+
cuda-flags = [
|
| 26 |
+
"-res-usage",
|
| 27 |
+
"--use_fast_math",
|
| 28 |
+
"-O3",
|
| 29 |
+
"--extra-device-vectorization",
|
| 30 |
+
"-DTmax=1024",
|
| 31 |
+
]
|
build/torch210-cu128-x86_64-windows/__init__.py
ADDED
|
@@ -0,0 +1,170 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from ._ops import ops
|
| 2 |
+
from typing import Tuple, Any
|
| 3 |
+
|
| 4 |
+
# Use a broad Tensor alias to avoid importing torch at import time.
|
| 5 |
+
from torch import Tensor
|
| 6 |
+
|
| 7 |
+
def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 8 |
+
"""RWKV WKV forward pass (float32).
|
| 9 |
+
|
| 10 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 14 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 15 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 16 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 17 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 18 |
+
|
| 19 |
+
Notes:
|
| 20 |
+
- All tensors must be on the same CUDA device.
|
| 21 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 22 |
+
"""
|
| 23 |
+
_validate_device_match((w, u, k, v, y))
|
| 24 |
+
ops.forward(w, u, k, v, y)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 28 |
+
"""RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).
|
| 29 |
+
|
| 30 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 31 |
+
|
| 32 |
+
Args:
|
| 33 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 34 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 35 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 36 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 37 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 38 |
+
|
| 39 |
+
Notes:
|
| 40 |
+
- All tensors must be on the same CUDA device.
|
| 41 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 42 |
+
"""
|
| 43 |
+
_validate_device_match((w, u, k, v, y))
|
| 44 |
+
ops.forward_bf16(w, u, k, v, y)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 48 |
+
"""RWKV WKV forward pass with persistent state (float32).
|
| 49 |
+
|
| 50 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 54 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 55 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 56 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 57 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 58 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 59 |
+
|
| 60 |
+
Notes:
|
| 61 |
+
- All tensors must be on the same CUDA device.
|
| 62 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 63 |
+
"""
|
| 64 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 65 |
+
ops.forward_with_state(w, u, k, v, y, s)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 69 |
+
"""RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).
|
| 70 |
+
|
| 71 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 72 |
+
|
| 73 |
+
Args:
|
| 74 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 75 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 76 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 77 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 78 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 79 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 80 |
+
|
| 81 |
+
Notes:
|
| 82 |
+
- All tensors must be on the same CUDA device.
|
| 83 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 84 |
+
"""
|
| 85 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 86 |
+
ops.forward_with_state_bf16(w, u, k, v, y, s)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def backward(
|
| 90 |
+
w: Tensor,
|
| 91 |
+
u: Tensor,
|
| 92 |
+
k: Tensor,
|
| 93 |
+
v: Tensor,
|
| 94 |
+
y: Tensor,
|
| 95 |
+
gy: Tensor,
|
| 96 |
+
gw: Tensor,
|
| 97 |
+
gu: Tensor,
|
| 98 |
+
gk: Tensor,
|
| 99 |
+
gv: Tensor,
|
| 100 |
+
) -> None:
|
| 101 |
+
"""RWKV WKV backward pass (float32).
|
| 102 |
+
|
| 103 |
+
Writes gradients into the provided tensors in-place.
|
| 104 |
+
|
| 105 |
+
Args:
|
| 106 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 107 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 108 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 109 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).
|
| 110 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 111 |
+
|
| 112 |
+
Notes:
|
| 113 |
+
- All tensors must be on the same CUDA device.
|
| 114 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 115 |
+
"""
|
| 116 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 117 |
+
ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def backward_bf16(
|
| 121 |
+
w: Tensor,
|
| 122 |
+
u: Tensor,
|
| 123 |
+
k: Tensor,
|
| 124 |
+
v: Tensor,
|
| 125 |
+
y: Tensor,
|
| 126 |
+
gy: Tensor,
|
| 127 |
+
gw: Tensor,
|
| 128 |
+
gu: Tensor,
|
| 129 |
+
gk: Tensor,
|
| 130 |
+
gv: Tensor,
|
| 131 |
+
) -> None:
|
| 132 |
+
"""RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).
|
| 133 |
+
|
| 134 |
+
Writes gradients into the provided tensors in-place.
|
| 135 |
+
|
| 136 |
+
Args:
|
| 137 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 138 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 139 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 140 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 141 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 142 |
+
|
| 143 |
+
Notes:
|
| 144 |
+
- All tensors must be on the same CUDA device.
|
| 145 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 146 |
+
"""
|
| 147 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 148 |
+
ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None:
|
| 152 |
+
"""Minimal runtime validation that all tensors live on the same CUDA device."""
|
| 153 |
+
if not tensors:
|
| 154 |
+
return
|
| 155 |
+
device = tensors[0].device
|
| 156 |
+
if not device.type == "cuda":
|
| 157 |
+
raise RuntimeError("RWKV CUDA ops require CUDA tensors")
|
| 158 |
+
for t in tensors[1:]:
|
| 159 |
+
if t.device != device:
|
| 160 |
+
raise RuntimeError("All tensors must be on the same CUDA device")
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
__all__ = [
|
| 164 |
+
"forward",
|
| 165 |
+
"forward_bf16",
|
| 166 |
+
"forward_with_state",
|
| 167 |
+
"forward_with_state_bf16",
|
| 168 |
+
"backward",
|
| 169 |
+
"backward_bf16",
|
| 170 |
+
]
|
build/torch210-cu128-x86_64-windows/_ops.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from . import _rwkv_cuda_38ccc47
|
| 3 |
+
ops = torch.ops._rwkv_cuda_38ccc47
|
| 4 |
+
|
| 5 |
+
def add_op_namespace_prefix(op_name: str):
|
| 6 |
+
"""
|
| 7 |
+
Prefix op by namespace.
|
| 8 |
+
"""
|
| 9 |
+
return f"_rwkv_cuda_38ccc47::{op_name}"
|
build/torch210-cu128-x86_64-windows/_rwkv_cuda_38ccc47.pyd
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e29ecb0b1d346a26decac4671633cf56fbc0b12df10027b1cca06a41f48976a7
|
| 3 |
+
size 423424
|
build/torch210-cu128-x86_64-windows/metadata.json
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"version": 1,
|
| 3 |
+
"license": "Apache-2.0",
|
| 4 |
+
"python-depends": [],
|
| 5 |
+
"backend": {
|
| 6 |
+
"type": "cuda",
|
| 7 |
+
"archs": [
|
| 8 |
+
"10.0",
|
| 9 |
+
"12.0",
|
| 10 |
+
"8.0",
|
| 11 |
+
"8.9",
|
| 12 |
+
"9.0"
|
| 13 |
+
]
|
| 14 |
+
}
|
| 15 |
+
}
|
build/torch210-cu128-x86_64-windows/rwkv/__init__.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import ctypes
|
| 2 |
+
import sys
|
| 3 |
+
|
| 4 |
+
import importlib
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
from types import ModuleType
|
| 7 |
+
|
| 8 |
+
def _import_from_path(file_path: Path) -> ModuleType:
|
| 9 |
+
# We cannot use the module name as-is, after adding it to `sys.modules`,
|
| 10 |
+
# it would also be used for other imports. So, we make a module name that
|
| 11 |
+
# depends on the path for it to be unique using the hex-encoded hash of
|
| 12 |
+
# the path.
|
| 13 |
+
path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
|
| 14 |
+
module_name = path_hash
|
| 15 |
+
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
| 16 |
+
if spec is None:
|
| 17 |
+
raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
|
| 18 |
+
module = importlib.util.module_from_spec(spec)
|
| 19 |
+
if module is None:
|
| 20 |
+
raise ImportError(f"Cannot load module {module_name} from spec")
|
| 21 |
+
sys.modules[module_name] = module
|
| 22 |
+
spec.loader.exec_module(module) # type: ignore
|
| 23 |
+
return module
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
|
build/torch210-cxx11-cu126-aarch64-linux/__init__.py
ADDED
|
@@ -0,0 +1,170 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from ._ops import ops
|
| 2 |
+
from typing import Tuple, Any
|
| 3 |
+
|
| 4 |
+
# Use a broad Tensor alias to avoid importing torch at import time.
|
| 5 |
+
from torch import Tensor
|
| 6 |
+
|
| 7 |
+
def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 8 |
+
"""RWKV WKV forward pass (float32).
|
| 9 |
+
|
| 10 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 14 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 15 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 16 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 17 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 18 |
+
|
| 19 |
+
Notes:
|
| 20 |
+
- All tensors must be on the same CUDA device.
|
| 21 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 22 |
+
"""
|
| 23 |
+
_validate_device_match((w, u, k, v, y))
|
| 24 |
+
ops.forward(w, u, k, v, y)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 28 |
+
"""RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).
|
| 29 |
+
|
| 30 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 31 |
+
|
| 32 |
+
Args:
|
| 33 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 34 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 35 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 36 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 37 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 38 |
+
|
| 39 |
+
Notes:
|
| 40 |
+
- All tensors must be on the same CUDA device.
|
| 41 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 42 |
+
"""
|
| 43 |
+
_validate_device_match((w, u, k, v, y))
|
| 44 |
+
ops.forward_bf16(w, u, k, v, y)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 48 |
+
"""RWKV WKV forward pass with persistent state (float32).
|
| 49 |
+
|
| 50 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 54 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 55 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 56 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 57 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 58 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 59 |
+
|
| 60 |
+
Notes:
|
| 61 |
+
- All tensors must be on the same CUDA device.
|
| 62 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 63 |
+
"""
|
| 64 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 65 |
+
ops.forward_with_state(w, u, k, v, y, s)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 69 |
+
"""RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).
|
| 70 |
+
|
| 71 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 72 |
+
|
| 73 |
+
Args:
|
| 74 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 75 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 76 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 77 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 78 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 79 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 80 |
+
|
| 81 |
+
Notes:
|
| 82 |
+
- All tensors must be on the same CUDA device.
|
| 83 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 84 |
+
"""
|
| 85 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 86 |
+
ops.forward_with_state_bf16(w, u, k, v, y, s)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def backward(
|
| 90 |
+
w: Tensor,
|
| 91 |
+
u: Tensor,
|
| 92 |
+
k: Tensor,
|
| 93 |
+
v: Tensor,
|
| 94 |
+
y: Tensor,
|
| 95 |
+
gy: Tensor,
|
| 96 |
+
gw: Tensor,
|
| 97 |
+
gu: Tensor,
|
| 98 |
+
gk: Tensor,
|
| 99 |
+
gv: Tensor,
|
| 100 |
+
) -> None:
|
| 101 |
+
"""RWKV WKV backward pass (float32).
|
| 102 |
+
|
| 103 |
+
Writes gradients into the provided tensors in-place.
|
| 104 |
+
|
| 105 |
+
Args:
|
| 106 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 107 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 108 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 109 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).
|
| 110 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 111 |
+
|
| 112 |
+
Notes:
|
| 113 |
+
- All tensors must be on the same CUDA device.
|
| 114 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 115 |
+
"""
|
| 116 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 117 |
+
ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def backward_bf16(
|
| 121 |
+
w: Tensor,
|
| 122 |
+
u: Tensor,
|
| 123 |
+
k: Tensor,
|
| 124 |
+
v: Tensor,
|
| 125 |
+
y: Tensor,
|
| 126 |
+
gy: Tensor,
|
| 127 |
+
gw: Tensor,
|
| 128 |
+
gu: Tensor,
|
| 129 |
+
gk: Tensor,
|
| 130 |
+
gv: Tensor,
|
| 131 |
+
) -> None:
|
| 132 |
+
"""RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).
|
| 133 |
+
|
| 134 |
+
Writes gradients into the provided tensors in-place.
|
| 135 |
+
|
| 136 |
+
Args:
|
| 137 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 138 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 139 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 140 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 141 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 142 |
+
|
| 143 |
+
Notes:
|
| 144 |
+
- All tensors must be on the same CUDA device.
|
| 145 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 146 |
+
"""
|
| 147 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 148 |
+
ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None:
|
| 152 |
+
"""Minimal runtime validation that all tensors live on the same CUDA device."""
|
| 153 |
+
if not tensors:
|
| 154 |
+
return
|
| 155 |
+
device = tensors[0].device
|
| 156 |
+
if not device.type == "cuda":
|
| 157 |
+
raise RuntimeError("RWKV CUDA ops require CUDA tensors")
|
| 158 |
+
for t in tensors[1:]:
|
| 159 |
+
if t.device != device:
|
| 160 |
+
raise RuntimeError("All tensors must be on the same CUDA device")
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
__all__ = [
|
| 164 |
+
"forward",
|
| 165 |
+
"forward_bf16",
|
| 166 |
+
"forward_with_state",
|
| 167 |
+
"forward_with_state_bf16",
|
| 168 |
+
"backward",
|
| 169 |
+
"backward_bf16",
|
| 170 |
+
]
|
build/torch210-cxx11-cu126-aarch64-linux/_ops.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from . import _rwkv_cuda_5849bdb
|
| 3 |
+
ops = torch.ops._rwkv_cuda_5849bdb
|
| 4 |
+
|
| 5 |
+
def add_op_namespace_prefix(op_name: str):
|
| 6 |
+
"""
|
| 7 |
+
Prefix op by namespace.
|
| 8 |
+
"""
|
| 9 |
+
return f"_rwkv_cuda_5849bdb::{op_name}"
|
build/torch210-cxx11-cu126-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:6eedf562b917d7e6865bbeb76296d462c5dd519365ab0454bc31558788c889c3
|
| 3 |
+
size 2232072
|
build/torch210-cxx11-cu126-aarch64-linux/metadata.json
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"version": 1,
|
| 3 |
+
"license": "Apache-2.0",
|
| 4 |
+
"python-depends": [],
|
| 5 |
+
"backend": {
|
| 6 |
+
"type": "cuda",
|
| 7 |
+
"archs": [
|
| 8 |
+
"8.0",
|
| 9 |
+
"8.9",
|
| 10 |
+
"9.0"
|
| 11 |
+
]
|
| 12 |
+
}
|
| 13 |
+
}
|
build/torch210-cxx11-cu126-aarch64-linux/rwkv/__init__.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import ctypes
|
| 2 |
+
import importlib.util
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from types import ModuleType
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def _import_from_path(file_path: Path) -> ModuleType:
|
| 9 |
+
# We cannot use the module name as-is, after adding it to `sys.modules`,
|
| 10 |
+
# it would also be used for other imports. So, we make a module name that
|
| 11 |
+
# depends on the path for it to be unique using the hex-encoded hash of
|
| 12 |
+
# the path.
|
| 13 |
+
path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
|
| 14 |
+
module_name = path_hash
|
| 15 |
+
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
| 16 |
+
if spec is None:
|
| 17 |
+
raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
|
| 18 |
+
module = importlib.util.module_from_spec(spec)
|
| 19 |
+
if module is None:
|
| 20 |
+
raise ImportError(f"Cannot load module {module_name} from spec")
|
| 21 |
+
sys.modules[module_name] = module
|
| 22 |
+
spec.loader.exec_module(module) # type: ignore
|
| 23 |
+
return module
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
|
build/torch210-cxx11-cu126-x86_64-linux/__init__.py
ADDED
|
@@ -0,0 +1,170 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from ._ops import ops
|
| 2 |
+
from typing import Tuple, Any
|
| 3 |
+
|
| 4 |
+
# Use a broad Tensor alias to avoid importing torch at import time.
|
| 5 |
+
from torch import Tensor
|
| 6 |
+
|
| 7 |
+
def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 8 |
+
"""RWKV WKV forward pass (float32).
|
| 9 |
+
|
| 10 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 14 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 15 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 16 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 17 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 18 |
+
|
| 19 |
+
Notes:
|
| 20 |
+
- All tensors must be on the same CUDA device.
|
| 21 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 22 |
+
"""
|
| 23 |
+
_validate_device_match((w, u, k, v, y))
|
| 24 |
+
ops.forward(w, u, k, v, y)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 28 |
+
"""RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).
|
| 29 |
+
|
| 30 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 31 |
+
|
| 32 |
+
Args:
|
| 33 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 34 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 35 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 36 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 37 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 38 |
+
|
| 39 |
+
Notes:
|
| 40 |
+
- All tensors must be on the same CUDA device.
|
| 41 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 42 |
+
"""
|
| 43 |
+
_validate_device_match((w, u, k, v, y))
|
| 44 |
+
ops.forward_bf16(w, u, k, v, y)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 48 |
+
"""RWKV WKV forward pass with persistent state (float32).
|
| 49 |
+
|
| 50 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 54 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 55 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 56 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 57 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 58 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 59 |
+
|
| 60 |
+
Notes:
|
| 61 |
+
- All tensors must be on the same CUDA device.
|
| 62 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 63 |
+
"""
|
| 64 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 65 |
+
ops.forward_with_state(w, u, k, v, y, s)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 69 |
+
"""RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).
|
| 70 |
+
|
| 71 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 72 |
+
|
| 73 |
+
Args:
|
| 74 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 75 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 76 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 77 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 78 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 79 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 80 |
+
|
| 81 |
+
Notes:
|
| 82 |
+
- All tensors must be on the same CUDA device.
|
| 83 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 84 |
+
"""
|
| 85 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 86 |
+
ops.forward_with_state_bf16(w, u, k, v, y, s)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def backward(
|
| 90 |
+
w: Tensor,
|
| 91 |
+
u: Tensor,
|
| 92 |
+
k: Tensor,
|
| 93 |
+
v: Tensor,
|
| 94 |
+
y: Tensor,
|
| 95 |
+
gy: Tensor,
|
| 96 |
+
gw: Tensor,
|
| 97 |
+
gu: Tensor,
|
| 98 |
+
gk: Tensor,
|
| 99 |
+
gv: Tensor,
|
| 100 |
+
) -> None:
|
| 101 |
+
"""RWKV WKV backward pass (float32).
|
| 102 |
+
|
| 103 |
+
Writes gradients into the provided tensors in-place.
|
| 104 |
+
|
| 105 |
+
Args:
|
| 106 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 107 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 108 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 109 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).
|
| 110 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 111 |
+
|
| 112 |
+
Notes:
|
| 113 |
+
- All tensors must be on the same CUDA device.
|
| 114 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 115 |
+
"""
|
| 116 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 117 |
+
ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def backward_bf16(
|
| 121 |
+
w: Tensor,
|
| 122 |
+
u: Tensor,
|
| 123 |
+
k: Tensor,
|
| 124 |
+
v: Tensor,
|
| 125 |
+
y: Tensor,
|
| 126 |
+
gy: Tensor,
|
| 127 |
+
gw: Tensor,
|
| 128 |
+
gu: Tensor,
|
| 129 |
+
gk: Tensor,
|
| 130 |
+
gv: Tensor,
|
| 131 |
+
) -> None:
|
| 132 |
+
"""RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).
|
| 133 |
+
|
| 134 |
+
Writes gradients into the provided tensors in-place.
|
| 135 |
+
|
| 136 |
+
Args:
|
| 137 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 138 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 139 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 140 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 141 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 142 |
+
|
| 143 |
+
Notes:
|
| 144 |
+
- All tensors must be on the same CUDA device.
|
| 145 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 146 |
+
"""
|
| 147 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 148 |
+
ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None:
|
| 152 |
+
"""Minimal runtime validation that all tensors live on the same CUDA device."""
|
| 153 |
+
if not tensors:
|
| 154 |
+
return
|
| 155 |
+
device = tensors[0].device
|
| 156 |
+
if not device.type == "cuda":
|
| 157 |
+
raise RuntimeError("RWKV CUDA ops require CUDA tensors")
|
| 158 |
+
for t in tensors[1:]:
|
| 159 |
+
if t.device != device:
|
| 160 |
+
raise RuntimeError("All tensors must be on the same CUDA device")
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
__all__ = [
|
| 164 |
+
"forward",
|
| 165 |
+
"forward_bf16",
|
| 166 |
+
"forward_with_state",
|
| 167 |
+
"forward_with_state_bf16",
|
| 168 |
+
"backward",
|
| 169 |
+
"backward_bf16",
|
| 170 |
+
]
|
build/torch210-cxx11-cu126-x86_64-linux/_ops.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from . import _rwkv_cuda_5849bdb
|
| 3 |
+
ops = torch.ops._rwkv_cuda_5849bdb
|
| 4 |
+
|
| 5 |
+
def add_op_namespace_prefix(op_name: str):
|
| 6 |
+
"""
|
| 7 |
+
Prefix op by namespace.
|
| 8 |
+
"""
|
| 9 |
+
return f"_rwkv_cuda_5849bdb::{op_name}"
|
build/torch210-cxx11-cu126-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:ef57a3b7b3028cf74874ad5f99fd237a942d20322b6c47e824cdcde75a612ac7
|
| 3 |
+
size 2116408
|
build/torch210-cxx11-cu126-x86_64-linux/metadata.json
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"version": 1,
|
| 3 |
+
"license": "Apache-2.0",
|
| 4 |
+
"python-depends": [],
|
| 5 |
+
"backend": {
|
| 6 |
+
"type": "cuda",
|
| 7 |
+
"archs": [
|
| 8 |
+
"8.0",
|
| 9 |
+
"8.9",
|
| 10 |
+
"9.0"
|
| 11 |
+
]
|
| 12 |
+
}
|
| 13 |
+
}
|
build/torch210-cxx11-cu126-x86_64-linux/rwkv/__init__.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import ctypes
|
| 2 |
+
import importlib.util
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from types import ModuleType
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def _import_from_path(file_path: Path) -> ModuleType:
|
| 9 |
+
# We cannot use the module name as-is, after adding it to `sys.modules`,
|
| 10 |
+
# it would also be used for other imports. So, we make a module name that
|
| 11 |
+
# depends on the path for it to be unique using the hex-encoded hash of
|
| 12 |
+
# the path.
|
| 13 |
+
path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
|
| 14 |
+
module_name = path_hash
|
| 15 |
+
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
| 16 |
+
if spec is None:
|
| 17 |
+
raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
|
| 18 |
+
module = importlib.util.module_from_spec(spec)
|
| 19 |
+
if module is None:
|
| 20 |
+
raise ImportError(f"Cannot load module {module_name} from spec")
|
| 21 |
+
sys.modules[module_name] = module
|
| 22 |
+
spec.loader.exec_module(module) # type: ignore
|
| 23 |
+
return module
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
|
build/torch210-cxx11-cu128-aarch64-linux/__init__.py
ADDED
|
@@ -0,0 +1,170 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from ._ops import ops
|
| 2 |
+
from typing import Tuple, Any
|
| 3 |
+
|
| 4 |
+
# Use a broad Tensor alias to avoid importing torch at import time.
|
| 5 |
+
from torch import Tensor
|
| 6 |
+
|
| 7 |
+
def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 8 |
+
"""RWKV WKV forward pass (float32).
|
| 9 |
+
|
| 10 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 14 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 15 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 16 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 17 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 18 |
+
|
| 19 |
+
Notes:
|
| 20 |
+
- All tensors must be on the same CUDA device.
|
| 21 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 22 |
+
"""
|
| 23 |
+
_validate_device_match((w, u, k, v, y))
|
| 24 |
+
ops.forward(w, u, k, v, y)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 28 |
+
"""RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).
|
| 29 |
+
|
| 30 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 31 |
+
|
| 32 |
+
Args:
|
| 33 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 34 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 35 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 36 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 37 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 38 |
+
|
| 39 |
+
Notes:
|
| 40 |
+
- All tensors must be on the same CUDA device.
|
| 41 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 42 |
+
"""
|
| 43 |
+
_validate_device_match((w, u, k, v, y))
|
| 44 |
+
ops.forward_bf16(w, u, k, v, y)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 48 |
+
"""RWKV WKV forward pass with persistent state (float32).
|
| 49 |
+
|
| 50 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 54 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 55 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 56 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 57 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 58 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 59 |
+
|
| 60 |
+
Notes:
|
| 61 |
+
- All tensors must be on the same CUDA device.
|
| 62 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 63 |
+
"""
|
| 64 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 65 |
+
ops.forward_with_state(w, u, k, v, y, s)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 69 |
+
"""RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).
|
| 70 |
+
|
| 71 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 72 |
+
|
| 73 |
+
Args:
|
| 74 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 75 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 76 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 77 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 78 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 79 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 80 |
+
|
| 81 |
+
Notes:
|
| 82 |
+
- All tensors must be on the same CUDA device.
|
| 83 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 84 |
+
"""
|
| 85 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 86 |
+
ops.forward_with_state_bf16(w, u, k, v, y, s)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def backward(
|
| 90 |
+
w: Tensor,
|
| 91 |
+
u: Tensor,
|
| 92 |
+
k: Tensor,
|
| 93 |
+
v: Tensor,
|
| 94 |
+
y: Tensor,
|
| 95 |
+
gy: Tensor,
|
| 96 |
+
gw: Tensor,
|
| 97 |
+
gu: Tensor,
|
| 98 |
+
gk: Tensor,
|
| 99 |
+
gv: Tensor,
|
| 100 |
+
) -> None:
|
| 101 |
+
"""RWKV WKV backward pass (float32).
|
| 102 |
+
|
| 103 |
+
Writes gradients into the provided tensors in-place.
|
| 104 |
+
|
| 105 |
+
Args:
|
| 106 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 107 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 108 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 109 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).
|
| 110 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 111 |
+
|
| 112 |
+
Notes:
|
| 113 |
+
- All tensors must be on the same CUDA device.
|
| 114 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 115 |
+
"""
|
| 116 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 117 |
+
ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def backward_bf16(
|
| 121 |
+
w: Tensor,
|
| 122 |
+
u: Tensor,
|
| 123 |
+
k: Tensor,
|
| 124 |
+
v: Tensor,
|
| 125 |
+
y: Tensor,
|
| 126 |
+
gy: Tensor,
|
| 127 |
+
gw: Tensor,
|
| 128 |
+
gu: Tensor,
|
| 129 |
+
gk: Tensor,
|
| 130 |
+
gv: Tensor,
|
| 131 |
+
) -> None:
|
| 132 |
+
"""RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).
|
| 133 |
+
|
| 134 |
+
Writes gradients into the provided tensors in-place.
|
| 135 |
+
|
| 136 |
+
Args:
|
| 137 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 138 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 139 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 140 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 141 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 142 |
+
|
| 143 |
+
Notes:
|
| 144 |
+
- All tensors must be on the same CUDA device.
|
| 145 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 146 |
+
"""
|
| 147 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 148 |
+
ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None:
|
| 152 |
+
"""Minimal runtime validation that all tensors live on the same CUDA device."""
|
| 153 |
+
if not tensors:
|
| 154 |
+
return
|
| 155 |
+
device = tensors[0].device
|
| 156 |
+
if not device.type == "cuda":
|
| 157 |
+
raise RuntimeError("RWKV CUDA ops require CUDA tensors")
|
| 158 |
+
for t in tensors[1:]:
|
| 159 |
+
if t.device != device:
|
| 160 |
+
raise RuntimeError("All tensors must be on the same CUDA device")
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
__all__ = [
|
| 164 |
+
"forward",
|
| 165 |
+
"forward_bf16",
|
| 166 |
+
"forward_with_state",
|
| 167 |
+
"forward_with_state_bf16",
|
| 168 |
+
"backward",
|
| 169 |
+
"backward_bf16",
|
| 170 |
+
]
|
build/torch210-cxx11-cu128-aarch64-linux/_ops.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from . import _rwkv_cuda_5849bdb
|
| 3 |
+
ops = torch.ops._rwkv_cuda_5849bdb
|
| 4 |
+
|
| 5 |
+
def add_op_namespace_prefix(op_name: str):
|
| 6 |
+
"""
|
| 7 |
+
Prefix op by namespace.
|
| 8 |
+
"""
|
| 9 |
+
return f"_rwkv_cuda_5849bdb::{op_name}"
|
build/torch210-cxx11-cu128-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:7baac5a668392dd6029e0b87b9d54f163e0c41ca6052a35a5d03116773784114
|
| 3 |
+
size 2428792
|
build/torch210-cxx11-cu128-aarch64-linux/metadata.json
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"version": 1,
|
| 3 |
+
"license": "Apache-2.0",
|
| 4 |
+
"python-depends": [],
|
| 5 |
+
"backend": {
|
| 6 |
+
"type": "cuda",
|
| 7 |
+
"archs": [
|
| 8 |
+
"10.0",
|
| 9 |
+
"12.0",
|
| 10 |
+
"8.0",
|
| 11 |
+
"8.9",
|
| 12 |
+
"9.0"
|
| 13 |
+
]
|
| 14 |
+
}
|
| 15 |
+
}
|
build/torch210-cxx11-cu128-aarch64-linux/rwkv/__init__.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import ctypes
|
| 2 |
+
import importlib.util
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from types import ModuleType
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def _import_from_path(file_path: Path) -> ModuleType:
|
| 9 |
+
# We cannot use the module name as-is, after adding it to `sys.modules`,
|
| 10 |
+
# it would also be used for other imports. So, we make a module name that
|
| 11 |
+
# depends on the path for it to be unique using the hex-encoded hash of
|
| 12 |
+
# the path.
|
| 13 |
+
path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
|
| 14 |
+
module_name = path_hash
|
| 15 |
+
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
| 16 |
+
if spec is None:
|
| 17 |
+
raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
|
| 18 |
+
module = importlib.util.module_from_spec(spec)
|
| 19 |
+
if module is None:
|
| 20 |
+
raise ImportError(f"Cannot load module {module_name} from spec")
|
| 21 |
+
sys.modules[module_name] = module
|
| 22 |
+
spec.loader.exec_module(module) # type: ignore
|
| 23 |
+
return module
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
|
build/torch210-cxx11-cu128-x86_64-linux/__init__.py
ADDED
|
@@ -0,0 +1,170 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from ._ops import ops
|
| 2 |
+
from typing import Tuple, Any
|
| 3 |
+
|
| 4 |
+
# Use a broad Tensor alias to avoid importing torch at import time.
|
| 5 |
+
from torch import Tensor
|
| 6 |
+
|
| 7 |
+
def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 8 |
+
"""RWKV WKV forward pass (float32).
|
| 9 |
+
|
| 10 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 14 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 15 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 16 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 17 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 18 |
+
|
| 19 |
+
Notes:
|
| 20 |
+
- All tensors must be on the same CUDA device.
|
| 21 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 22 |
+
"""
|
| 23 |
+
_validate_device_match((w, u, k, v, y))
|
| 24 |
+
ops.forward(w, u, k, v, y)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 28 |
+
"""RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).
|
| 29 |
+
|
| 30 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 31 |
+
|
| 32 |
+
Args:
|
| 33 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 34 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 35 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 36 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 37 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 38 |
+
|
| 39 |
+
Notes:
|
| 40 |
+
- All tensors must be on the same CUDA device.
|
| 41 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 42 |
+
"""
|
| 43 |
+
_validate_device_match((w, u, k, v, y))
|
| 44 |
+
ops.forward_bf16(w, u, k, v, y)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 48 |
+
"""RWKV WKV forward pass with persistent state (float32).
|
| 49 |
+
|
| 50 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 54 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 55 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 56 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 57 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 58 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 59 |
+
|
| 60 |
+
Notes:
|
| 61 |
+
- All tensors must be on the same CUDA device.
|
| 62 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 63 |
+
"""
|
| 64 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 65 |
+
ops.forward_with_state(w, u, k, v, y, s)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 69 |
+
"""RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).
|
| 70 |
+
|
| 71 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 72 |
+
|
| 73 |
+
Args:
|
| 74 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 75 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 76 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 77 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 78 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 79 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 80 |
+
|
| 81 |
+
Notes:
|
| 82 |
+
- All tensors must be on the same CUDA device.
|
| 83 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 84 |
+
"""
|
| 85 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 86 |
+
ops.forward_with_state_bf16(w, u, k, v, y, s)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def backward(
|
| 90 |
+
w: Tensor,
|
| 91 |
+
u: Tensor,
|
| 92 |
+
k: Tensor,
|
| 93 |
+
v: Tensor,
|
| 94 |
+
y: Tensor,
|
| 95 |
+
gy: Tensor,
|
| 96 |
+
gw: Tensor,
|
| 97 |
+
gu: Tensor,
|
| 98 |
+
gk: Tensor,
|
| 99 |
+
gv: Tensor,
|
| 100 |
+
) -> None:
|
| 101 |
+
"""RWKV WKV backward pass (float32).
|
| 102 |
+
|
| 103 |
+
Writes gradients into the provided tensors in-place.
|
| 104 |
+
|
| 105 |
+
Args:
|
| 106 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 107 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 108 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 109 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).
|
| 110 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 111 |
+
|
| 112 |
+
Notes:
|
| 113 |
+
- All tensors must be on the same CUDA device.
|
| 114 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 115 |
+
"""
|
| 116 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 117 |
+
ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def backward_bf16(
|
| 121 |
+
w: Tensor,
|
| 122 |
+
u: Tensor,
|
| 123 |
+
k: Tensor,
|
| 124 |
+
v: Tensor,
|
| 125 |
+
y: Tensor,
|
| 126 |
+
gy: Tensor,
|
| 127 |
+
gw: Tensor,
|
| 128 |
+
gu: Tensor,
|
| 129 |
+
gk: Tensor,
|
| 130 |
+
gv: Tensor,
|
| 131 |
+
) -> None:
|
| 132 |
+
"""RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).
|
| 133 |
+
|
| 134 |
+
Writes gradients into the provided tensors in-place.
|
| 135 |
+
|
| 136 |
+
Args:
|
| 137 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 138 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 139 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 140 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 141 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 142 |
+
|
| 143 |
+
Notes:
|
| 144 |
+
- All tensors must be on the same CUDA device.
|
| 145 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 146 |
+
"""
|
| 147 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 148 |
+
ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None:
|
| 152 |
+
"""Minimal runtime validation that all tensors live on the same CUDA device."""
|
| 153 |
+
if not tensors:
|
| 154 |
+
return
|
| 155 |
+
device = tensors[0].device
|
| 156 |
+
if not device.type == "cuda":
|
| 157 |
+
raise RuntimeError("RWKV CUDA ops require CUDA tensors")
|
| 158 |
+
for t in tensors[1:]:
|
| 159 |
+
if t.device != device:
|
| 160 |
+
raise RuntimeError("All tensors must be on the same CUDA device")
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
__all__ = [
|
| 164 |
+
"forward",
|
| 165 |
+
"forward_bf16",
|
| 166 |
+
"forward_with_state",
|
| 167 |
+
"forward_with_state_bf16",
|
| 168 |
+
"backward",
|
| 169 |
+
"backward_bf16",
|
| 170 |
+
]
|
build/torch210-cxx11-cu128-x86_64-linux/_ops.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from . import _rwkv_cuda_5849bdb
|
| 3 |
+
ops = torch.ops._rwkv_cuda_5849bdb
|
| 4 |
+
|
| 5 |
+
def add_op_namespace_prefix(op_name: str):
|
| 6 |
+
"""
|
| 7 |
+
Prefix op by namespace.
|
| 8 |
+
"""
|
| 9 |
+
return f"_rwkv_cuda_5849bdb::{op_name}"
|
build/torch210-cxx11-cu128-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:bb33f213f9412b8120ed540d92ef1e9c70ed35590feff5d53113140ce5298ee0
|
| 3 |
+
size 2318768
|
build/torch210-cxx11-cu128-x86_64-linux/metadata.json
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"version": 1,
|
| 3 |
+
"license": "Apache-2.0",
|
| 4 |
+
"python-depends": [],
|
| 5 |
+
"backend": {
|
| 6 |
+
"type": "cuda",
|
| 7 |
+
"archs": [
|
| 8 |
+
"10.0",
|
| 9 |
+
"12.0",
|
| 10 |
+
"8.0",
|
| 11 |
+
"8.9",
|
| 12 |
+
"9.0"
|
| 13 |
+
]
|
| 14 |
+
}
|
| 15 |
+
}
|
build/torch210-cxx11-cu128-x86_64-linux/rwkv/__init__.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import ctypes
|
| 2 |
+
import importlib.util
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from types import ModuleType
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def _import_from_path(file_path: Path) -> ModuleType:
|
| 9 |
+
# We cannot use the module name as-is, after adding it to `sys.modules`,
|
| 10 |
+
# it would also be used for other imports. So, we make a module name that
|
| 11 |
+
# depends on the path for it to be unique using the hex-encoded hash of
|
| 12 |
+
# the path.
|
| 13 |
+
path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
|
| 14 |
+
module_name = path_hash
|
| 15 |
+
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
| 16 |
+
if spec is None:
|
| 17 |
+
raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
|
| 18 |
+
module = importlib.util.module_from_spec(spec)
|
| 19 |
+
if module is None:
|
| 20 |
+
raise ImportError(f"Cannot load module {module_name} from spec")
|
| 21 |
+
sys.modules[module_name] = module
|
| 22 |
+
spec.loader.exec_module(module) # type: ignore
|
| 23 |
+
return module
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
|
build/torch210-cxx11-cu130-aarch64-linux/__init__.py
ADDED
|
@@ -0,0 +1,170 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from ._ops import ops
|
| 2 |
+
from typing import Tuple, Any
|
| 3 |
+
|
| 4 |
+
# Use a broad Tensor alias to avoid importing torch at import time.
|
| 5 |
+
from torch import Tensor
|
| 6 |
+
|
| 7 |
+
def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 8 |
+
"""RWKV WKV forward pass (float32).
|
| 9 |
+
|
| 10 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 14 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 15 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 16 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 17 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 18 |
+
|
| 19 |
+
Notes:
|
| 20 |
+
- All tensors must be on the same CUDA device.
|
| 21 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 22 |
+
"""
|
| 23 |
+
_validate_device_match((w, u, k, v, y))
|
| 24 |
+
ops.forward(w, u, k, v, y)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 28 |
+
"""RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).
|
| 29 |
+
|
| 30 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 31 |
+
|
| 32 |
+
Args:
|
| 33 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 34 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 35 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 36 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 37 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 38 |
+
|
| 39 |
+
Notes:
|
| 40 |
+
- All tensors must be on the same CUDA device.
|
| 41 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 42 |
+
"""
|
| 43 |
+
_validate_device_match((w, u, k, v, y))
|
| 44 |
+
ops.forward_bf16(w, u, k, v, y)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 48 |
+
"""RWKV WKV forward pass with persistent state (float32).
|
| 49 |
+
|
| 50 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 54 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 55 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 56 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 57 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 58 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 59 |
+
|
| 60 |
+
Notes:
|
| 61 |
+
- All tensors must be on the same CUDA device.
|
| 62 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 63 |
+
"""
|
| 64 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 65 |
+
ops.forward_with_state(w, u, k, v, y, s)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 69 |
+
"""RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).
|
| 70 |
+
|
| 71 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 72 |
+
|
| 73 |
+
Args:
|
| 74 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 75 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 76 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 77 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 78 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 79 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 80 |
+
|
| 81 |
+
Notes:
|
| 82 |
+
- All tensors must be on the same CUDA device.
|
| 83 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 84 |
+
"""
|
| 85 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 86 |
+
ops.forward_with_state_bf16(w, u, k, v, y, s)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def backward(
|
| 90 |
+
w: Tensor,
|
| 91 |
+
u: Tensor,
|
| 92 |
+
k: Tensor,
|
| 93 |
+
v: Tensor,
|
| 94 |
+
y: Tensor,
|
| 95 |
+
gy: Tensor,
|
| 96 |
+
gw: Tensor,
|
| 97 |
+
gu: Tensor,
|
| 98 |
+
gk: Tensor,
|
| 99 |
+
gv: Tensor,
|
| 100 |
+
) -> None:
|
| 101 |
+
"""RWKV WKV backward pass (float32).
|
| 102 |
+
|
| 103 |
+
Writes gradients into the provided tensors in-place.
|
| 104 |
+
|
| 105 |
+
Args:
|
| 106 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 107 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 108 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 109 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).
|
| 110 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 111 |
+
|
| 112 |
+
Notes:
|
| 113 |
+
- All tensors must be on the same CUDA device.
|
| 114 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 115 |
+
"""
|
| 116 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 117 |
+
ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def backward_bf16(
|
| 121 |
+
w: Tensor,
|
| 122 |
+
u: Tensor,
|
| 123 |
+
k: Tensor,
|
| 124 |
+
v: Tensor,
|
| 125 |
+
y: Tensor,
|
| 126 |
+
gy: Tensor,
|
| 127 |
+
gw: Tensor,
|
| 128 |
+
gu: Tensor,
|
| 129 |
+
gk: Tensor,
|
| 130 |
+
gv: Tensor,
|
| 131 |
+
) -> None:
|
| 132 |
+
"""RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).
|
| 133 |
+
|
| 134 |
+
Writes gradients into the provided tensors in-place.
|
| 135 |
+
|
| 136 |
+
Args:
|
| 137 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 138 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 139 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 140 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 141 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 142 |
+
|
| 143 |
+
Notes:
|
| 144 |
+
- All tensors must be on the same CUDA device.
|
| 145 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 146 |
+
"""
|
| 147 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 148 |
+
ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None:
|
| 152 |
+
"""Minimal runtime validation that all tensors live on the same CUDA device."""
|
| 153 |
+
if not tensors:
|
| 154 |
+
return
|
| 155 |
+
device = tensors[0].device
|
| 156 |
+
if not device.type == "cuda":
|
| 157 |
+
raise RuntimeError("RWKV CUDA ops require CUDA tensors")
|
| 158 |
+
for t in tensors[1:]:
|
| 159 |
+
if t.device != device:
|
| 160 |
+
raise RuntimeError("All tensors must be on the same CUDA device")
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
__all__ = [
|
| 164 |
+
"forward",
|
| 165 |
+
"forward_bf16",
|
| 166 |
+
"forward_with_state",
|
| 167 |
+
"forward_with_state_bf16",
|
| 168 |
+
"backward",
|
| 169 |
+
"backward_bf16",
|
| 170 |
+
]
|
build/torch210-cxx11-cu130-aarch64-linux/_ops.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from . import _rwkv_cuda_5849bdb
|
| 3 |
+
ops = torch.ops._rwkv_cuda_5849bdb
|
| 4 |
+
|
| 5 |
+
def add_op_namespace_prefix(op_name: str):
|
| 6 |
+
"""
|
| 7 |
+
Prefix op by namespace.
|
| 8 |
+
"""
|
| 9 |
+
return f"_rwkv_cuda_5849bdb::{op_name}"
|
build/torch210-cxx11-cu130-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:2a90cc433fa1453bb04a4993f2af0485266c227d0ece13808be55d3a7f280f4f
|
| 3 |
+
size 2432776
|
build/torch210-cxx11-cu130-aarch64-linux/metadata.json
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"version": 1,
|
| 3 |
+
"license": "Apache-2.0",
|
| 4 |
+
"python-depends": [],
|
| 5 |
+
"backend": {
|
| 6 |
+
"type": "cuda",
|
| 7 |
+
"archs": [
|
| 8 |
+
"10.0",
|
| 9 |
+
"12.0",
|
| 10 |
+
"8.0",
|
| 11 |
+
"8.9",
|
| 12 |
+
"9.0"
|
| 13 |
+
]
|
| 14 |
+
}
|
| 15 |
+
}
|
build/torch210-cxx11-cu130-aarch64-linux/rwkv/__init__.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import ctypes
|
| 2 |
+
import importlib.util
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from types import ModuleType
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def _import_from_path(file_path: Path) -> ModuleType:
|
| 9 |
+
# We cannot use the module name as-is, after adding it to `sys.modules`,
|
| 10 |
+
# it would also be used for other imports. So, we make a module name that
|
| 11 |
+
# depends on the path for it to be unique using the hex-encoded hash of
|
| 12 |
+
# the path.
|
| 13 |
+
path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
|
| 14 |
+
module_name = path_hash
|
| 15 |
+
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
| 16 |
+
if spec is None:
|
| 17 |
+
raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
|
| 18 |
+
module = importlib.util.module_from_spec(spec)
|
| 19 |
+
if module is None:
|
| 20 |
+
raise ImportError(f"Cannot load module {module_name} from spec")
|
| 21 |
+
sys.modules[module_name] = module
|
| 22 |
+
spec.loader.exec_module(module) # type: ignore
|
| 23 |
+
return module
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
|
build/torch210-cxx11-cu130-x86_64-linux/__init__.py
ADDED
|
@@ -0,0 +1,170 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from ._ops import ops
|
| 2 |
+
from typing import Tuple, Any
|
| 3 |
+
|
| 4 |
+
# Use a broad Tensor alias to avoid importing torch at import time.
|
| 5 |
+
from torch import Tensor
|
| 6 |
+
|
| 7 |
+
def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 8 |
+
"""RWKV WKV forward pass (float32).
|
| 9 |
+
|
| 10 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 14 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 15 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 16 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 17 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 18 |
+
|
| 19 |
+
Notes:
|
| 20 |
+
- All tensors must be on the same CUDA device.
|
| 21 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 22 |
+
"""
|
| 23 |
+
_validate_device_match((w, u, k, v, y))
|
| 24 |
+
ops.forward(w, u, k, v, y)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 28 |
+
"""RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).
|
| 29 |
+
|
| 30 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 31 |
+
|
| 32 |
+
Args:
|
| 33 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 34 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 35 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 36 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 37 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 38 |
+
|
| 39 |
+
Notes:
|
| 40 |
+
- All tensors must be on the same CUDA device.
|
| 41 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 42 |
+
"""
|
| 43 |
+
_validate_device_match((w, u, k, v, y))
|
| 44 |
+
ops.forward_bf16(w, u, k, v, y)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 48 |
+
"""RWKV WKV forward pass with persistent state (float32).
|
| 49 |
+
|
| 50 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 54 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 55 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 56 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 57 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 58 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 59 |
+
|
| 60 |
+
Notes:
|
| 61 |
+
- All tensors must be on the same CUDA device.
|
| 62 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 63 |
+
"""
|
| 64 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 65 |
+
ops.forward_with_state(w, u, k, v, y, s)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 69 |
+
"""RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).
|
| 70 |
+
|
| 71 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 72 |
+
|
| 73 |
+
Args:
|
| 74 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 75 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 76 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 77 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 78 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 79 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 80 |
+
|
| 81 |
+
Notes:
|
| 82 |
+
- All tensors must be on the same CUDA device.
|
| 83 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 84 |
+
"""
|
| 85 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 86 |
+
ops.forward_with_state_bf16(w, u, k, v, y, s)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def backward(
|
| 90 |
+
w: Tensor,
|
| 91 |
+
u: Tensor,
|
| 92 |
+
k: Tensor,
|
| 93 |
+
v: Tensor,
|
| 94 |
+
y: Tensor,
|
| 95 |
+
gy: Tensor,
|
| 96 |
+
gw: Tensor,
|
| 97 |
+
gu: Tensor,
|
| 98 |
+
gk: Tensor,
|
| 99 |
+
gv: Tensor,
|
| 100 |
+
) -> None:
|
| 101 |
+
"""RWKV WKV backward pass (float32).
|
| 102 |
+
|
| 103 |
+
Writes gradients into the provided tensors in-place.
|
| 104 |
+
|
| 105 |
+
Args:
|
| 106 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 107 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 108 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 109 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).
|
| 110 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 111 |
+
|
| 112 |
+
Notes:
|
| 113 |
+
- All tensors must be on the same CUDA device.
|
| 114 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 115 |
+
"""
|
| 116 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 117 |
+
ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def backward_bf16(
|
| 121 |
+
w: Tensor,
|
| 122 |
+
u: Tensor,
|
| 123 |
+
k: Tensor,
|
| 124 |
+
v: Tensor,
|
| 125 |
+
y: Tensor,
|
| 126 |
+
gy: Tensor,
|
| 127 |
+
gw: Tensor,
|
| 128 |
+
gu: Tensor,
|
| 129 |
+
gk: Tensor,
|
| 130 |
+
gv: Tensor,
|
| 131 |
+
) -> None:
|
| 132 |
+
"""RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).
|
| 133 |
+
|
| 134 |
+
Writes gradients into the provided tensors in-place.
|
| 135 |
+
|
| 136 |
+
Args:
|
| 137 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 138 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 139 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 140 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 141 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 142 |
+
|
| 143 |
+
Notes:
|
| 144 |
+
- All tensors must be on the same CUDA device.
|
| 145 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 146 |
+
"""
|
| 147 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 148 |
+
ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None:
|
| 152 |
+
"""Minimal runtime validation that all tensors live on the same CUDA device."""
|
| 153 |
+
if not tensors:
|
| 154 |
+
return
|
| 155 |
+
device = tensors[0].device
|
| 156 |
+
if not device.type == "cuda":
|
| 157 |
+
raise RuntimeError("RWKV CUDA ops require CUDA tensors")
|
| 158 |
+
for t in tensors[1:]:
|
| 159 |
+
if t.device != device:
|
| 160 |
+
raise RuntimeError("All tensors must be on the same CUDA device")
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
__all__ = [
|
| 164 |
+
"forward",
|
| 165 |
+
"forward_bf16",
|
| 166 |
+
"forward_with_state",
|
| 167 |
+
"forward_with_state_bf16",
|
| 168 |
+
"backward",
|
| 169 |
+
"backward_bf16",
|
| 170 |
+
]
|
build/torch210-cxx11-cu130-x86_64-linux/_ops.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from . import _rwkv_cuda_5849bdb
|
| 3 |
+
ops = torch.ops._rwkv_cuda_5849bdb
|
| 4 |
+
|
| 5 |
+
def add_op_namespace_prefix(op_name: str):
|
| 6 |
+
"""
|
| 7 |
+
Prefix op by namespace.
|
| 8 |
+
"""
|
| 9 |
+
return f"_rwkv_cuda_5849bdb::{op_name}"
|
build/torch210-cxx11-cu130-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:024c9733ab553416a3165fabcfd1f8957e99b2c256c3fc5f51d1e3032200ff96
|
| 3 |
+
size 2348248
|
build/torch210-cxx11-cu130-x86_64-linux/metadata.json
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"version": 1,
|
| 3 |
+
"license": "Apache-2.0",
|
| 4 |
+
"python-depends": [],
|
| 5 |
+
"backend": {
|
| 6 |
+
"type": "cuda",
|
| 7 |
+
"archs": [
|
| 8 |
+
"10.0",
|
| 9 |
+
"12.0",
|
| 10 |
+
"8.0",
|
| 11 |
+
"8.9",
|
| 12 |
+
"9.0"
|
| 13 |
+
]
|
| 14 |
+
}
|
| 15 |
+
}
|
build/torch210-cxx11-cu130-x86_64-linux/rwkv/__init__.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import ctypes
|
| 2 |
+
import importlib.util
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from types import ModuleType
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def _import_from_path(file_path: Path) -> ModuleType:
|
| 9 |
+
# We cannot use the module name as-is, after adding it to `sys.modules`,
|
| 10 |
+
# it would also be used for other imports. So, we make a module name that
|
| 11 |
+
# depends on the path for it to be unique using the hex-encoded hash of
|
| 12 |
+
# the path.
|
| 13 |
+
path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
|
| 14 |
+
module_name = path_hash
|
| 15 |
+
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
| 16 |
+
if spec is None:
|
| 17 |
+
raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
|
| 18 |
+
module = importlib.util.module_from_spec(spec)
|
| 19 |
+
if module is None:
|
| 20 |
+
raise ImportError(f"Cannot load module {module_name} from spec")
|
| 21 |
+
sys.modules[module_name] = module
|
| 22 |
+
spec.loader.exec_module(module) # type: ignore
|
| 23 |
+
return module
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
|
build/torch211-cxx11-cu126-aarch64-linux/__init__.py
ADDED
|
@@ -0,0 +1,170 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from ._ops import ops
|
| 2 |
+
from typing import Tuple, Any
|
| 3 |
+
|
| 4 |
+
# Use a broad Tensor alias to avoid importing torch at import time.
|
| 5 |
+
from torch import Tensor
|
| 6 |
+
|
| 7 |
+
def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 8 |
+
"""RWKV WKV forward pass (float32).
|
| 9 |
+
|
| 10 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 14 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 15 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 16 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 17 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 18 |
+
|
| 19 |
+
Notes:
|
| 20 |
+
- All tensors must be on the same CUDA device.
|
| 21 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 22 |
+
"""
|
| 23 |
+
_validate_device_match((w, u, k, v, y))
|
| 24 |
+
ops.forward(w, u, k, v, y)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 28 |
+
"""RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).
|
| 29 |
+
|
| 30 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 31 |
+
|
| 32 |
+
Args:
|
| 33 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 34 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 35 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 36 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 37 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 38 |
+
|
| 39 |
+
Notes:
|
| 40 |
+
- All tensors must be on the same CUDA device.
|
| 41 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 42 |
+
"""
|
| 43 |
+
_validate_device_match((w, u, k, v, y))
|
| 44 |
+
ops.forward_bf16(w, u, k, v, y)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 48 |
+
"""RWKV WKV forward pass with persistent state (float32).
|
| 49 |
+
|
| 50 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 54 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 55 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 56 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 57 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 58 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 59 |
+
|
| 60 |
+
Notes:
|
| 61 |
+
- All tensors must be on the same CUDA device.
|
| 62 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 63 |
+
"""
|
| 64 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 65 |
+
ops.forward_with_state(w, u, k, v, y, s)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 69 |
+
"""RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).
|
| 70 |
+
|
| 71 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 72 |
+
|
| 73 |
+
Args:
|
| 74 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 75 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 76 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 77 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 78 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 79 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 80 |
+
|
| 81 |
+
Notes:
|
| 82 |
+
- All tensors must be on the same CUDA device.
|
| 83 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 84 |
+
"""
|
| 85 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 86 |
+
ops.forward_with_state_bf16(w, u, k, v, y, s)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def backward(
|
| 90 |
+
w: Tensor,
|
| 91 |
+
u: Tensor,
|
| 92 |
+
k: Tensor,
|
| 93 |
+
v: Tensor,
|
| 94 |
+
y: Tensor,
|
| 95 |
+
gy: Tensor,
|
| 96 |
+
gw: Tensor,
|
| 97 |
+
gu: Tensor,
|
| 98 |
+
gk: Tensor,
|
| 99 |
+
gv: Tensor,
|
| 100 |
+
) -> None:
|
| 101 |
+
"""RWKV WKV backward pass (float32).
|
| 102 |
+
|
| 103 |
+
Writes gradients into the provided tensors in-place.
|
| 104 |
+
|
| 105 |
+
Args:
|
| 106 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 107 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 108 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 109 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).
|
| 110 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 111 |
+
|
| 112 |
+
Notes:
|
| 113 |
+
- All tensors must be on the same CUDA device.
|
| 114 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 115 |
+
"""
|
| 116 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 117 |
+
ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def backward_bf16(
|
| 121 |
+
w: Tensor,
|
| 122 |
+
u: Tensor,
|
| 123 |
+
k: Tensor,
|
| 124 |
+
v: Tensor,
|
| 125 |
+
y: Tensor,
|
| 126 |
+
gy: Tensor,
|
| 127 |
+
gw: Tensor,
|
| 128 |
+
gu: Tensor,
|
| 129 |
+
gk: Tensor,
|
| 130 |
+
gv: Tensor,
|
| 131 |
+
) -> None:
|
| 132 |
+
"""RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).
|
| 133 |
+
|
| 134 |
+
Writes gradients into the provided tensors in-place.
|
| 135 |
+
|
| 136 |
+
Args:
|
| 137 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 138 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 139 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 140 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 141 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 142 |
+
|
| 143 |
+
Notes:
|
| 144 |
+
- All tensors must be on the same CUDA device.
|
| 145 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 146 |
+
"""
|
| 147 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 148 |
+
ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None:
|
| 152 |
+
"""Minimal runtime validation that all tensors live on the same CUDA device."""
|
| 153 |
+
if not tensors:
|
| 154 |
+
return
|
| 155 |
+
device = tensors[0].device
|
| 156 |
+
if not device.type == "cuda":
|
| 157 |
+
raise RuntimeError("RWKV CUDA ops require CUDA tensors")
|
| 158 |
+
for t in tensors[1:]:
|
| 159 |
+
if t.device != device:
|
| 160 |
+
raise RuntimeError("All tensors must be on the same CUDA device")
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
__all__ = [
|
| 164 |
+
"forward",
|
| 165 |
+
"forward_bf16",
|
| 166 |
+
"forward_with_state",
|
| 167 |
+
"forward_with_state_bf16",
|
| 168 |
+
"backward",
|
| 169 |
+
"backward_bf16",
|
| 170 |
+
]
|
build/torch211-cxx11-cu126-aarch64-linux/_ops.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from . import _rwkv_cuda_5849bdb
|
| 3 |
+
ops = torch.ops._rwkv_cuda_5849bdb
|
| 4 |
+
|
| 5 |
+
def add_op_namespace_prefix(op_name: str):
|
| 6 |
+
"""
|
| 7 |
+
Prefix op by namespace.
|
| 8 |
+
"""
|
| 9 |
+
return f"_rwkv_cuda_5849bdb::{op_name}"
|
build/torch211-cxx11-cu126-aarch64-linux/_rwkv_cuda_5849bdb.abi3.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:f97c2ec3958d76db37b74a83093e2f127317b51c94d089bc0ac710499fbdc49a
|
| 3 |
+
size 2232072
|
build/torch211-cxx11-cu126-aarch64-linux/metadata.json
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"version": 1,
|
| 3 |
+
"license": "Apache-2.0",
|
| 4 |
+
"python-depends": [],
|
| 5 |
+
"backend": {
|
| 6 |
+
"type": "cuda",
|
| 7 |
+
"archs": [
|
| 8 |
+
"8.0",
|
| 9 |
+
"8.9",
|
| 10 |
+
"9.0"
|
| 11 |
+
]
|
| 12 |
+
}
|
| 13 |
+
}
|
build/torch211-cxx11-cu126-aarch64-linux/rwkv/__init__.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import ctypes
|
| 2 |
+
import importlib.util
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from types import ModuleType
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def _import_from_path(file_path: Path) -> ModuleType:
|
| 9 |
+
# We cannot use the module name as-is, after adding it to `sys.modules`,
|
| 10 |
+
# it would also be used for other imports. So, we make a module name that
|
| 11 |
+
# depends on the path for it to be unique using the hex-encoded hash of
|
| 12 |
+
# the path.
|
| 13 |
+
path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
|
| 14 |
+
module_name = path_hash
|
| 15 |
+
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
| 16 |
+
if spec is None:
|
| 17 |
+
raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
|
| 18 |
+
module = importlib.util.module_from_spec(spec)
|
| 19 |
+
if module is None:
|
| 20 |
+
raise ImportError(f"Cannot load module {module_name} from spec")
|
| 21 |
+
sys.modules[module_name] = module
|
| 22 |
+
spec.loader.exec_module(module) # type: ignore
|
| 23 |
+
return module
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
|
build/torch211-cxx11-cu126-x86_64-linux/__init__.py
ADDED
|
@@ -0,0 +1,170 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from ._ops import ops
|
| 2 |
+
from typing import Tuple, Any
|
| 3 |
+
|
| 4 |
+
# Use a broad Tensor alias to avoid importing torch at import time.
|
| 5 |
+
from torch import Tensor
|
| 6 |
+
|
| 7 |
+
def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 8 |
+
"""RWKV WKV forward pass (float32).
|
| 9 |
+
|
| 10 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 14 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 15 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 16 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 17 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 18 |
+
|
| 19 |
+
Notes:
|
| 20 |
+
- All tensors must be on the same CUDA device.
|
| 21 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 22 |
+
"""
|
| 23 |
+
_validate_device_match((w, u, k, v, y))
|
| 24 |
+
ops.forward(w, u, k, v, y)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 28 |
+
"""RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).
|
| 29 |
+
|
| 30 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 31 |
+
|
| 32 |
+
Args:
|
| 33 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 34 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 35 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 36 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 37 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 38 |
+
|
| 39 |
+
Notes:
|
| 40 |
+
- All tensors must be on the same CUDA device.
|
| 41 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 42 |
+
"""
|
| 43 |
+
_validate_device_match((w, u, k, v, y))
|
| 44 |
+
ops.forward_bf16(w, u, k, v, y)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 48 |
+
"""RWKV WKV forward pass with persistent state (float32).
|
| 49 |
+
|
| 50 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 54 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 55 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 56 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 57 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 58 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 59 |
+
|
| 60 |
+
Notes:
|
| 61 |
+
- All tensors must be on the same CUDA device.
|
| 62 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 63 |
+
"""
|
| 64 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 65 |
+
ops.forward_with_state(w, u, k, v, y, s)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 69 |
+
"""RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).
|
| 70 |
+
|
| 71 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 72 |
+
|
| 73 |
+
Args:
|
| 74 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 75 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 76 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 77 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 78 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 79 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 80 |
+
|
| 81 |
+
Notes:
|
| 82 |
+
- All tensors must be on the same CUDA device.
|
| 83 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 84 |
+
"""
|
| 85 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 86 |
+
ops.forward_with_state_bf16(w, u, k, v, y, s)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def backward(
|
| 90 |
+
w: Tensor,
|
| 91 |
+
u: Tensor,
|
| 92 |
+
k: Tensor,
|
| 93 |
+
v: Tensor,
|
| 94 |
+
y: Tensor,
|
| 95 |
+
gy: Tensor,
|
| 96 |
+
gw: Tensor,
|
| 97 |
+
gu: Tensor,
|
| 98 |
+
gk: Tensor,
|
| 99 |
+
gv: Tensor,
|
| 100 |
+
) -> None:
|
| 101 |
+
"""RWKV WKV backward pass (float32).
|
| 102 |
+
|
| 103 |
+
Writes gradients into the provided tensors in-place.
|
| 104 |
+
|
| 105 |
+
Args:
|
| 106 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 107 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 108 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 109 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).
|
| 110 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 111 |
+
|
| 112 |
+
Notes:
|
| 113 |
+
- All tensors must be on the same CUDA device.
|
| 114 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 115 |
+
"""
|
| 116 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 117 |
+
ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def backward_bf16(
|
| 121 |
+
w: Tensor,
|
| 122 |
+
u: Tensor,
|
| 123 |
+
k: Tensor,
|
| 124 |
+
v: Tensor,
|
| 125 |
+
y: Tensor,
|
| 126 |
+
gy: Tensor,
|
| 127 |
+
gw: Tensor,
|
| 128 |
+
gu: Tensor,
|
| 129 |
+
gk: Tensor,
|
| 130 |
+
gv: Tensor,
|
| 131 |
+
) -> None:
|
| 132 |
+
"""RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).
|
| 133 |
+
|
| 134 |
+
Writes gradients into the provided tensors in-place.
|
| 135 |
+
|
| 136 |
+
Args:
|
| 137 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 138 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 139 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 140 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 141 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 142 |
+
|
| 143 |
+
Notes:
|
| 144 |
+
- All tensors must be on the same CUDA device.
|
| 145 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 146 |
+
"""
|
| 147 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 148 |
+
ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None:
|
| 152 |
+
"""Minimal runtime validation that all tensors live on the same CUDA device."""
|
| 153 |
+
if not tensors:
|
| 154 |
+
return
|
| 155 |
+
device = tensors[0].device
|
| 156 |
+
if not device.type == "cuda":
|
| 157 |
+
raise RuntimeError("RWKV CUDA ops require CUDA tensors")
|
| 158 |
+
for t in tensors[1:]:
|
| 159 |
+
if t.device != device:
|
| 160 |
+
raise RuntimeError("All tensors must be on the same CUDA device")
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
__all__ = [
|
| 164 |
+
"forward",
|
| 165 |
+
"forward_bf16",
|
| 166 |
+
"forward_with_state",
|
| 167 |
+
"forward_with_state_bf16",
|
| 168 |
+
"backward",
|
| 169 |
+
"backward_bf16",
|
| 170 |
+
]
|
build/torch211-cxx11-cu126-x86_64-linux/_ops.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from . import _rwkv_cuda_5849bdb
|
| 3 |
+
ops = torch.ops._rwkv_cuda_5849bdb
|
| 4 |
+
|
| 5 |
+
def add_op_namespace_prefix(op_name: str):
|
| 6 |
+
"""
|
| 7 |
+
Prefix op by namespace.
|
| 8 |
+
"""
|
| 9 |
+
return f"_rwkv_cuda_5849bdb::{op_name}"
|
build/torch211-cxx11-cu126-x86_64-linux/_rwkv_cuda_5849bdb.abi3.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a55661a0185ca9375700bbd80fe85272180c0f081e115d45ef53b471f55fc1e9
|
| 3 |
+
size 2116408
|
build/torch211-cxx11-cu126-x86_64-linux/metadata.json
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"version": 1,
|
| 3 |
+
"license": "Apache-2.0",
|
| 4 |
+
"python-depends": [],
|
| 5 |
+
"backend": {
|
| 6 |
+
"type": "cuda",
|
| 7 |
+
"archs": [
|
| 8 |
+
"8.0",
|
| 9 |
+
"8.9",
|
| 10 |
+
"9.0"
|
| 11 |
+
]
|
| 12 |
+
}
|
| 13 |
+
}
|
build/torch211-cxx11-cu126-x86_64-linux/rwkv/__init__.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import ctypes
|
| 2 |
+
import importlib.util
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from types import ModuleType
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def _import_from_path(file_path: Path) -> ModuleType:
|
| 9 |
+
# We cannot use the module name as-is, after adding it to `sys.modules`,
|
| 10 |
+
# it would also be used for other imports. So, we make a module name that
|
| 11 |
+
# depends on the path for it to be unique using the hex-encoded hash of
|
| 12 |
+
# the path.
|
| 13 |
+
path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
|
| 14 |
+
module_name = path_hash
|
| 15 |
+
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
| 16 |
+
if spec is None:
|
| 17 |
+
raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
|
| 18 |
+
module = importlib.util.module_from_spec(spec)
|
| 19 |
+
if module is None:
|
| 20 |
+
raise ImportError(f"Cannot load module {module_name} from spec")
|
| 21 |
+
sys.modules[module_name] = module
|
| 22 |
+
spec.loader.exec_module(module) # type: ignore
|
| 23 |
+
return module
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
|
build/torch211-cxx11-cu128-aarch64-linux/__init__.py
ADDED
|
@@ -0,0 +1,170 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from ._ops import ops
|
| 2 |
+
from typing import Tuple, Any
|
| 3 |
+
|
| 4 |
+
# Use a broad Tensor alias to avoid importing torch at import time.
|
| 5 |
+
from torch import Tensor
|
| 6 |
+
|
| 7 |
+
def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 8 |
+
"""RWKV WKV forward pass (float32).
|
| 9 |
+
|
| 10 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 14 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 15 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 16 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 17 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 18 |
+
|
| 19 |
+
Notes:
|
| 20 |
+
- All tensors must be on the same CUDA device.
|
| 21 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 22 |
+
"""
|
| 23 |
+
_validate_device_match((w, u, k, v, y))
|
| 24 |
+
ops.forward(w, u, k, v, y)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None:
|
| 28 |
+
"""RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).
|
| 29 |
+
|
| 30 |
+
Runs the CUDA kernel and writes the result into ``y`` in-place.
|
| 31 |
+
|
| 32 |
+
Args:
|
| 33 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 34 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 35 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 36 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 37 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 38 |
+
|
| 39 |
+
Notes:
|
| 40 |
+
- All tensors must be on the same CUDA device.
|
| 41 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 42 |
+
"""
|
| 43 |
+
_validate_device_match((w, u, k, v, y))
|
| 44 |
+
ops.forward_bf16(w, u, k, v, y)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 48 |
+
"""RWKV WKV forward pass with persistent state (float32).
|
| 49 |
+
|
| 50 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 54 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 55 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 56 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 57 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 58 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 59 |
+
|
| 60 |
+
Notes:
|
| 61 |
+
- All tensors must be on the same CUDA device.
|
| 62 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 63 |
+
"""
|
| 64 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 65 |
+
ops.forward_with_state(w, u, k, v, y, s)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None:
|
| 69 |
+
"""RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).
|
| 70 |
+
|
| 71 |
+
Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``.
|
| 72 |
+
|
| 73 |
+
Args:
|
| 74 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 75 |
+
u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 76 |
+
k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 77 |
+
v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 78 |
+
y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 79 |
+
s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).
|
| 80 |
+
|
| 81 |
+
Notes:
|
| 82 |
+
- All tensors must be on the same CUDA device.
|
| 83 |
+
- Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.
|
| 84 |
+
"""
|
| 85 |
+
_validate_device_match((w, u, k, v, y, s))
|
| 86 |
+
ops.forward_with_state_bf16(w, u, k, v, y, s)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def backward(
|
| 90 |
+
w: Tensor,
|
| 91 |
+
u: Tensor,
|
| 92 |
+
k: Tensor,
|
| 93 |
+
v: Tensor,
|
| 94 |
+
y: Tensor,
|
| 95 |
+
gy: Tensor,
|
| 96 |
+
gw: Tensor,
|
| 97 |
+
gu: Tensor,
|
| 98 |
+
gk: Tensor,
|
| 99 |
+
gv: Tensor,
|
| 100 |
+
) -> None:
|
| 101 |
+
"""RWKV WKV backward pass (float32).
|
| 102 |
+
|
| 103 |
+
Writes gradients into the provided tensors in-place.
|
| 104 |
+
|
| 105 |
+
Args:
|
| 106 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 107 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 108 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.
|
| 109 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).
|
| 110 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).
|
| 111 |
+
|
| 112 |
+
Notes:
|
| 113 |
+
- All tensors must be on the same CUDA device.
|
| 114 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 115 |
+
"""
|
| 116 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 117 |
+
ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def backward_bf16(
|
| 121 |
+
w: Tensor,
|
| 122 |
+
u: Tensor,
|
| 123 |
+
k: Tensor,
|
| 124 |
+
v: Tensor,
|
| 125 |
+
y: Tensor,
|
| 126 |
+
gy: Tensor,
|
| 127 |
+
gw: Tensor,
|
| 128 |
+
gu: Tensor,
|
| 129 |
+
gk: Tensor,
|
| 130 |
+
gv: Tensor,
|
| 131 |
+
) -> None:
|
| 132 |
+
"""RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).
|
| 133 |
+
|
| 134 |
+
Writes gradients into the provided tensors in-place.
|
| 135 |
+
|
| 136 |
+
Args:
|
| 137 |
+
w: Decay weights, shape ``[C]``, dtype ``torch.float32``.
|
| 138 |
+
u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 139 |
+
gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.
|
| 140 |
+
gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 141 |
+
gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).
|
| 142 |
+
|
| 143 |
+
Notes:
|
| 144 |
+
- All tensors must be on the same CUDA device.
|
| 145 |
+
- Shapes must agree on ``B``, ``T`` and ``C``.
|
| 146 |
+
"""
|
| 147 |
+
_validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv))
|
| 148 |
+
ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv)
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None:
|
| 152 |
+
"""Minimal runtime validation that all tensors live on the same CUDA device."""
|
| 153 |
+
if not tensors:
|
| 154 |
+
return
|
| 155 |
+
device = tensors[0].device
|
| 156 |
+
if not device.type == "cuda":
|
| 157 |
+
raise RuntimeError("RWKV CUDA ops require CUDA tensors")
|
| 158 |
+
for t in tensors[1:]:
|
| 159 |
+
if t.device != device:
|
| 160 |
+
raise RuntimeError("All tensors must be on the same CUDA device")
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
__all__ = [
|
| 164 |
+
"forward",
|
| 165 |
+
"forward_bf16",
|
| 166 |
+
"forward_with_state",
|
| 167 |
+
"forward_with_state_bf16",
|
| 168 |
+
"backward",
|
| 169 |
+
"backward_bf16",
|
| 170 |
+
]
|