rwkv-ops 0.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.

This version of rwkv-ops might be problematic. Click here for more details.

Files changed (7) hide show
  1. rwkv_ops/__init__.py +1 -1
  2. rwkv_ops/rwkv7_kernel/torch_op.py +24 -10
  3. {rwkv_ops-0.2.dist-info → rwkv_ops-0.2.1.dist-info}/METADATA +1 -1
  4. {rwkv_ops-0.2.dist-info → rwkv_ops-0.2.1.dist-info}/RECORD +7 -7
  5. {rwkv_ops-0.2.dist-info → rwkv_ops-0.2.1.dist-info}/LICENSE.txt +0 -0
  6. {rwkv_ops-0.2.dist-info → rwkv_ops-0.2.1.dist-info}/WHEEL +0 -0
  7. {rwkv_ops-0.2.dist-info → rwkv_ops-0.2.1.dist-info}/top_level.txt +0 -0
rwkv_ops/__init__.py CHANGED
@@ -1,4 +1,4 @@
1
- __version__ = "0.2"
1
+ __version__ = "0.2.1"
2
2
  import os
3
3
 
4
4
  KERNEL_TYPE = os.environ.get("KERNEL_TYPE", "triton")
rwkv_ops/rwkv7_kernel/torch_op.py CHANGED
@@ -473,16 +473,30 @@ def generalized_delta_rule(
473
473
  a = transpose_head(a, head_first)
474
474
  b = transpose_head(b, head_first)
475
475
  w = transpose_head(w, head_first)
476
- out, state = chunk_rwkv7(
477
- r=r,
478
- k=k,
479
- v=v,
480
- a=a,
481
- b=b,
482
- w=w,
483
- initial_state=initial_state,
484
- output_final_state=output_final_state,
485
- )
476
+ if w.device.type == "cuda":
477
+ out, state = chunk_rwkv7(
478
+ r=r,
479
+ k=k,
480
+ v=v,
481
+ a=a,
482
+ b=b,
483
+ w=w,
484
+ initial_state=initial_state,
485
+ output_final_state=output_final_state,
486
+ )
487
+ else:
488
+ from ops.native_keras_op import generalized_delta_rule
489
+
490
+ out, state = generalized_delta_rule(
491
+ r=r,
492
+ k=k,
493
+ v=v,
494
+ a=a,
495
+ b=b,
496
+ w=w,
497
+ initial_state=initial_state,
498
+ output_final_state=output_final_state,
499
+ )
486
500
  out = transpose_head(out, head_first)
487
501
  if output_final_state:
488
502
  return out, cast(state, dtype)
{rwkv_ops-0.2.dist-info → rwkv_ops-0.2.1.dist-info}/METADATA RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: rwkv-ops
3
- Version: 0.2
3
+ Version: 0.2.1
4
4
  Home-page: https://github.com/pass-lin/rwkv_ops
5
5
  License: Apache 2.0
6
6
  Keywords: rwkv implement for multi backend
{rwkv_ops-0.2.dist-info → rwkv_ops-0.2.1.dist-info}/RECORD RENAMED
@@ -1,4 +1,4 @@
1
- rwkv_ops/__init__.py,sha256=Kfw_9iearbIplpuxx8sUl30TxKbTO-Ehqe0290Y4sZw,841
1
+ rwkv_ops/__init__.py,sha256=FhqDH54a_vrLJQLTWMEM-QnETJG2z6l4YjMiLtcZt2E,843
2
2
  rwkv_ops/rwkv6_kernel/__init__.py,sha256=_j6G_3fY8xPxrlZbgDT2ndX4IPiNJ4qjqIcdmNI_r9Q,4100
3
3
  rwkv_ops/rwkv6_kernel/jax_rwkv_kernel.py,sha256=WOzqfQQSHHMoWqm2kRz_BhtMzGYc5USJ26qaEwuARo4,30117
4
4
  rwkv_ops/rwkv6_kernel/ops_rwkv_kernel.py,sha256=otjfw5n6nf2YVpBIWIZjaCsxMyLXXwg-ma1ueXX-EdY,3274
@@ -8,7 +8,7 @@ rwkv_ops/rwkv7_kernel/get_jax_devices_info.py,sha256=cMIaNED7d1PvYNSyq8wNI3G7wNv
8
8
  rwkv_ops/rwkv7_kernel/get_torch_devices_info.py,sha256=ZL_rAM6lHB4nTOOU28Xm08qptfuIoijOMi_xwJG3KCo,7380
9
9
  rwkv_ops/rwkv7_kernel/jax_op.py,sha256=tyMxvk_EblDaGsePpxw3AhELvolp7LeE5NopUhKw1R0,9107
10
10
  rwkv_ops/rwkv7_kernel/native_keras_op.py,sha256=QPrXLbqw0chipQg_0jepRp2U19BYpBBFdKZWyaDNNoc,2488
11
- rwkv_ops/rwkv7_kernel/torch_op.py,sha256=yY5QP87iDow-T6a4ZzFShyIQ8gprTQoLYcjFqgOTW4Y,13675
11
+ rwkv_ops/rwkv7_kernel/torch_op.py,sha256=jw_AvqshTAG4t9-MRqxFQNi_bTzxNbx3lwnMifPk8-8,14070
12
12
  rwkv_ops/rwkv7_kernel/jax_kernel/__init__.py,sha256=uHsf_1qrtRK62IvhLuzefHGPWpHXmw1p0tqmwlHcptk,346
13
13
  rwkv_ops/rwkv7_kernel/jax_kernel/chunk_A_bwd.py,sha256=2Voq1Bdzn0DFloiLvwINBk7akmxRWIqXIQeyafrJJGg,2138
14
14
  rwkv_ops/rwkv7_kernel/jax_kernel/chunk_A_fwd.py,sha256=rhmglqHIIww7yPzaSBEp9ISxhhxoUbMtV51AUDyhUd8,1425
@@ -40,8 +40,8 @@ rwkv_ops/rwkv7_kernel/triton_kernel/cumsum.py,sha256=pRp_z587PrnpgRVpi031IndyjVI
40
40
  rwkv_ops/rwkv7_kernel/triton_kernel/utils.py,sha256=TNGlkwGq4t-TOcdVBk_N_vHPLzMFTu_F0V-O1RprIO4,553
41
41
  rwkv_ops/rwkv7_kernel/triton_kernel/wy_fast_bwd.py,sha256=szaG11q_WmpyhXi6aVWwzizvflCh5wND8wGA_V8afzA,5479
42
42
  rwkv_ops/rwkv7_kernel/triton_kernel/wy_fast_fwd.py,sha256=jbb19DUTHENU2RIOv_T4m_W1eXMqdRqG0XevIkBOhI4,9438
43
- rwkv_ops-0.2.dist-info/LICENSE.txt,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
44
- rwkv_ops-0.2.dist-info/METADATA,sha256=cUhC6EYLULLgNLVtLOg4qMSoAZjIDcrojax6egCTU04,8409
45
- rwkv_ops-0.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
46
- rwkv_ops-0.2.dist-info/top_level.txt,sha256=cVqoKE-WR_e2gHL87-6O4K1kG6-yTJGB2huyr6FmD2I,9
47
- rwkv_ops-0.2.dist-info/RECORD,,
43
+ rwkv_ops-0.2.1.dist-info/LICENSE.txt,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
44
+ rwkv_ops-0.2.1.dist-info/METADATA,sha256=_I9thMqiJMM1kOlwUH23DVkRTnBHVLyjHQ5-74OvpS4,8411
45
+ rwkv_ops-0.2.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
46
+ rwkv_ops-0.2.1.dist-info/top_level.txt,sha256=cVqoKE-WR_e2gHL87-6O4K1kG6-yTJGB2huyr6FmD2I,9
47
+ rwkv_ops-0.2.1.dist-info/RECORD,,