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

rwkv_ops/rwkv7_kernel/torch_op.py

@@ -7,6 +7,7 @@ backward gradient computation, and integration with PyTorch's autograd system.
 该文件实现了分块 Delta Rule 注意力机制的前向与反向传播，
 使用 Triton 内核进行 GPU 加速优化。包括前向传播、梯度反向传播函数，
 并集成了 PyTorch 的自动求导系统。
+
 """
 
 import warnings
@@ -15,11 +16,11 @@ from typing import Optional
 import torch
 import triton
 
-# 导入内核实现模块 / Import kernel implementation modules
 from .torch_kernel.chunk_A_bwd import chunk_dplr_bwd_dqk_intra
 from .torch_kernel.chunk_A_fwd import chunk_dplr_fwd_intra
 from .torch_kernel.chunk_h_bwd import chunk_dplr_bwd_dhu
 from .torch_kernel.chunk_h_fwd import chunk_dplr_fwd_h
+
 from .torch_kernel.chunk_o_bwd import (
     chunk_dplr_bwd_dAu,
     chunk_dplr_bwd_dv,
@@ -37,11 +38,6 @@ from .get_torch_devices_info import (
 
 
 def cast(x, dtype):
-    """
-    Cast tensor x to specified dtype if not already in that format.
-
-    如果张量 x 不是目标数据类型，则将其转换为目标类型。
-    """
     if x is None or x.dtype == dtype:
         return x
     return x.to(dtype)
@@ -97,11 +93,11 @@ def chunk_dplr_fwd(
 
     del ge
 
-    # Compute WY representation
+    # A_ab, A_ak, gi, ge torch.float32
+    # A_qk, A_qb, qg, kg, ag, bg, dtype=q.dtype, eg: bf16
     w, u, _ = prepare_wy_repr_fwd(ag=ag, A_ab=A_ab, A_ak=A_ak, v=v, chunk_size=BT)
 
     del A_ab, A_ak
-
     h, v_new, final_state = chunk_dplr_fwd_h(
         kg=kg,
         bg=bg,
@@ -164,6 +160,7 @@ def chunk_dplr_bwd(
         dgk (torch.Tensor): Gradient of log decays
         dh0 (torch.Tensor): Gradient of initial state
     """
+    # ******* start recomputing everything, otherwise i believe the gpu memory will be exhausted *******
     gi, ge = chunk_rwkv6_fwd_cumsum(gk, BT)
     A_ab, A_qk, A_ak, A_qb, qg, kg, ag, bg = chunk_dplr_fwd_intra(
         q=q,
@@ -183,6 +180,9 @@ def chunk_dplr_bwd(
         kg=kg, bg=bg, v=v, w=w, u=u, gk=gi, initial_state=initial_state, chunk_size=BT
     )
     del u
+    # ******* end of recomputation *******
+    # A_ak, A_ab_inv, gi, ge torch.float32
+    # A_qk, A_qb, qg, kg, ag, bg, v_new dtype=q.dtype, eg: bf16
 
     dv_new_intra, dA_qk, dA_qb = chunk_dplr_bwd_dAu(
         v=v, v_new=v_new, do=do, A_qb=A_qb, scale=scale, chunk_size=BT
@@ -211,7 +211,7 @@ def chunk_dplr_bwd(
         do=do,
         h=h,
         dh=dh,
-        dv=dv,
+        dv=dv_new,
         w=w,
         gk=gi,
         chunk_size=BT,
@@ -282,11 +282,6 @@ class ChunkDPLRDeltaRuleFunction(torch.autograd.Function):
         output_final_state: bool = True,
         cu_seqlens: Optional[torch.LongTensor] = None,
     ):
-        """
-        Forward function with autograd support.
-
-        支持自动求导的前向函数。
-        """
         chunk_size = 16
         o, final_state = chunk_dplr_fwd(
             q=q,
@@ -310,11 +305,6 @@ class ChunkDPLRDeltaRuleFunction(torch.autograd.Function):
     @input_guard
     @autocast_custom_bwd
     def backward(ctx, do: torch.Tensor, dht: torch.Tensor):
-        """
-        Backward function with autograd support.
-
-        支持自动求导的反向函数。
-        """
         q, k, v, a, b, gk, initial_state = ctx.saved_tensors
         BT = ctx.chunk_size
         cu_seqlens = ctx.cu_seqlens
@@ -439,6 +429,7 @@ def chunk_rwkv7(
 
     RWKV-7 注意力机制的接口函数。
     """
+
     if w is not None:
         log_w = -torch.exp(w)
     else:
@@ -458,11 +449,6 @@ def chunk_rwkv7(
 
 
 def transpose_head(x, head_first):
-    """
-    Transpose between head-first and time-first formats.
-
-    在 head-first 和 time-first 格式之间转置。
-    """
     if head_first:
         x = torch.permute(x, dims=(0, 2, 1, 3))
     out = cast(x, torch.bfloat16).contiguous()
@@ -480,11 +466,6 @@ def generalized_delta_rule(
     output_final_state: bool = True,
     head_first: bool = False,
 ):
-    """
-    Generalized delta rule attention interface.
-
-    泛化 Delta Rule 注意力机制接口。
-    """
     dtype = r.dtype
     r = transpose_head(r, head_first)
     k = transpose_head(k, head_first)
@@ -492,30 +473,16 @@ def generalized_delta_rule(
     a = transpose_head(a, head_first)
     b = transpose_head(b, head_first)
     w = transpose_head(w, head_first)
-    if w.device.type == "cuda":
-        out, state = chunk_rwkv7(
-            r=r,
-            k=k,
-            v=v,
-            a=a,
-            b=b,
-            w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-        )
-    else:
-        from .native_keras_op import generalized_delta_rule
-
-        out, state = generalized_delta_rule(
-            r=r,
-            k=k,
-            v=v,
-            a=a,
-            b=b,
-            w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-        )
+    out, state = chunk_rwkv7(
+        r=r,
+        k=k,
+        v=v,
+        a=a,
+        b=b,
+        w=w,
+        initial_state=initial_state,
+        output_final_state=output_final_state,
+    )
     out = transpose_head(out, head_first)
     if output_final_state:
         return out, cast(state, dtype)

rwkv_ops-0.2.dist-info/METADATA

@@ -0,0 +1,258 @@
+Metadata-Version: 2.1
+Name: rwkv-ops
+Version: 0.2
+Home-page: https://github.com/pass-lin/rwkv_ops
+License: Apache 2.0
+Keywords: rwkv implement for multi backend
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Description-Content-Type: text/markdown
+License-File: LICENSE.txt
+Requires-Dist: keras
+
+[English Document](ENREADME.md)
+
+# RWKV OPS 项目
+
+> 由于 RWKV 将持续迭代，核心算子会随之更新。  
+> 本仓专门维护「算子」本身，不维护 layer 与 model；尽可能提供各框架的 GPU 算子。  
+
+### 当前支持
+| 算子类型 | 框架支持 |
+|----------|----------|
+| GPU 算子 | PyTorch、JAX（TensorFlow 待 Google 支持 Triton 后上线） |
+| 原生算子 | PyTorch、JAX、TensorFlow、NumPy |
+
+> 未来若 Keras 生态扩展，可能支持 MLX、OpenVINO。  
+> 注意：本库依赖 `keras`。
+
+---
+
+## 安装
+
+```bash
+pip install rwkv_ops
+```
+
+---
+
+## 环境变量
+
+| 变量名 | 含义 | 取值 | 默认值 | 优先级 |
+|---|---|---|---|---|
+| `KERAS_BACKEND` | Keras 后端 | `jax` / `torch` / `tensorflow` / `numpy` | — | 低 |
+| `KERNEL_BACKEND` | 算子后端 | `jax` / `torch` / `tensorflow` / `numpy` | `torch` | **高** |
+| `KERNEL_TYPE` | 实现类型 | `triton` / `cuda` / `native` | — | — |
+
+> 若 `KERNEL_BACKEND` 有值，直接采用；若为空，则用 `KERAS_BACKEND`；两者皆空则默认 `torch`。  
+> `native` 为原生算子，无 chunkwise，速度慢且显存高。
+
+---
+
+## rwkv7op 使用方法
+
+```python
+from rwkv_ops import generalized_delta_rule  # 或 from rwkv_ops import rwkv7_op，完全等价
+
+def generalized_delta_rule(
+    r,
+    w,
+    k,
+    v,
+    a,
+    b,
+    initial_state=None,
+    output_final_state: bool = True,
+    head_first: bool = False,
+):
+    """
+    分块 Delta Rule 注意力接口。
+
+    Args:
+        q:  [B, T, H, K]
+        k:  [B, T, H, K]
+        v:  [B, T, H, V]
+        a:  [B, T, H, K]
+        b:  [B, T, H, K]
+        gk: [B, T, H, K]  # decay term in log space!
+        initial_state: 初始状态 [N, H, K, V]，N 为序列数
+        output_final_state: 是否返回最终状态
+        head_first: 是否 head-first 格式，不支持变长
+
+    Returns:
+        o:           输出 [B, T, H, V] 或 [B, H, T, V]
+        final_state: 最终状态 [N, H, K, V] 或 None
+    """
+```
+
+### torch-cuda 特殊用法
+
+- torch-cuda 下 `head_size` 也是一个 kernel 参数，默认为 64。  
+- 若 `head_size ≠ 64`，请使用：
+
+```python
+from rwkv_ops import get_generalized_delta_rule
+
+generalized_delta_rule, RWKV7_USE_KERNEL = get_generalized_delta_rule(
+    your_head_size, KERNEL_TYPE="cuda"
+)
+```
+
+- `RWKV7_USE_KERNEL` 为常量，标记是否使用 chunkwise 算子。  
+- 两者 padding 处理逻辑不同：
+
+```python
+if padding_mask is not None:
+    if RWKV7_USE_KERNEL:
+        w += (1 - padding_mask) * -1e9
+    else:
+        w = w * padding_mask + 1 - padding_mask
+```
+
+---
+
+### rwkv7op 实现状态
+
+| Framework   | cuda | triton | native |
+|-------------|------|--------|--------|
+| PyTorch     | ✅   | ✅     | ✅     |
+| JAX         | ❌   | ✅     | ✅     |
+| TensorFlow  | ❌   | ❌     | ✅     |
+| NumPy       | ❌   | ❌     | ✅     |
+
+---
+
+## rwkv6op 使用方法
+
+### PyTorch 使用注意事项
+
+- 安装依赖：`keras`、`ninja`、完整的 CUDA 工具包。
+- 若使用 VS Code + 虚拟环境调试，请务必在终端手动激活虚拟环境，再运行代码，否则 ninja 可能无法工作。
+- 虽然 PyTorch 在「虚拟环境中的 CUDA 版本」与「全局 CUDA 版本」不一致时仍可正常运行，但强烈建议保持一致。
+- PyTorch 限制：同一程序内只能实例化 **一个** `RWKV6_OP` 对象；算子线程安全（无状态），可在多处调用。
+
+### JAX 使用注意事项
+
+- 安装依赖：`keras`、`gcc`、`pybind11`、完整的 CUDA 工具包。
+- 即使通过虚拟环境为 JAX 安装 CUDA，也必须在系统级安装完整 CUDA；两者版本需一致，以保证 JAX 并行编译速度。
+- JAX 编译依赖 `/usr/local/cuda` 软链接，如不存在请手动创建：
+  ```shell
+  sudo ln -sf /usr/local/cuda-12.4 /usr/local/cuda
+  ```
+- 确保 `nvcc -V` 正常输出，且 `which nvcc` 指向正确版本。
+- JAX 限制：同一程序内只能实例化 **一个** `RWKV6_OP` 对象；算子线程安全（无状态），可在多处调用。
+- JAX ≥ 0.6.0 不再使用 CUDA 算子，默认使用原生算子；推荐 0.4.34。
+
+### TensorFlow 使用注意事项
+
+- 仅提供基于原生 API 的 `RWKV6` 算子，仅用于推理，效率较低。
+
+---
+
+### 使用方法
+需要注意的是，和rwkv7写成函数的形式不一样，RWKV6的op是一个类，需要实例化。
+```python
+from rwkv_ops import RWKV6_OP
+
+operator = RWKV6_OP(
+    head_size=64,               # 头大小，不确定时填 64
+    max_sequence_length=4096,   # 训练最大序列长度；推理不受限
+    ops_loop=False              # 可选：序列长度=1 时是否用上层 API 替代 CUDA
+)
+```
+
+#### 调用
+
+```python
+y, y_state = operator(
+    r, k, v, w, u,
+    with_state=False,   # 是否使用自定义初始状态 / 输出结束状态
+    init_state=None,    # 初始状态 [n_state, num_heads, head_size, head_size]
+    state_map=None      # int32 一维数组，长度=batch_size，定义 init_state 映射
+)
+```
+
+| 参数 | 形状 | 说明 |
+|---|---|---|
+| r, k, v, w | (batch_size, seq_len, hidden_size) | — |
+| u | (num_heads, head_size) 或 (hidden_size,) | — |
+| init_state | (n_state, num_heads, head_size, head_size) | n_state=1 时所有样本共用；n_state=batch_size 时一一对应 |
+| state_map | (batch_size,) | 指定每个样本用到的 init_state 索引 |
+
+| 返回值 | 形状 | 说明 |
+|---|---|---|
+| y | (batch_size, seq_len, hidden_size) | 输出 |
+| y_state | (batch_size, num_heads, head_size, head_size) 或 None | 结束状态 |
+
+---
+
+### 分布式小贴士
+
+- 算子本身无分布式支持；PyTorch 可直接用多线程分布式。
+- JAX 需通过 `shard_map` 包装（示例）：
+
+```python
+import os
+os.environ['KERAS_BACKEND'] = 'jax'
+
+import jax, jax.numpy as jnp
+from jax.experimental.shard_map import shard_map
+from jax.sharding import Mesh, PartitionSpec as P
+from functools import partial
+from rwkv_ops import RWKV6_OP
+
+batch_size, seq_length = 24, 512
+head_size, num_heads = 64, 32
+hidden_size = head_size * num_heads
+
+mesh = Mesh(jax.devices('gpu'), axis_names=('device_axis',))
+device_ns = NamedSharding(mesh, P('device_axis'))
+
+operator = RWKV6_OP(head_size=head_size, max_sequence_length=seq_length)
+
+@partial(shard_map,
+         mesh=mesh,
+         in_specs=(P('device_axis'),) * 5,
+         out_specs=(P('device_axis'), P('device_axis')),
+         check_rep=False)
+def call_kernel(r, k, v, w, u):
+    # 去掉最外 device 维度
+    r, k, v, w, u = map(jnp.squeeze, (r, k, v, w, u))
+    y, ys = operator(r, k, v, w, u, with_state=True)
+    return jnp.expand_dims(y, 0), jnp.expand_dims(ys, 0)
+
+# 构造输入并放置到对应设备
+keys = jax.random.split(jax.random.PRNGKey(0), 5)
+inputs = [jax.random.normal(k, (mesh.size, batch_size, seq_length, hidden_size)) for k in keys]
+inputs_r, inputs_k, inputs_v, inputs_w, inputs_u = map(
+    lambda x: jax.device_put(x, device_ns), inputs)
+inputs_u = inputs_u[:, :, 0]  # (devices, hidden_size)
+
+# 可选：jax.jit(call_kernel, ...) 加速
+outputs_y, y_state = call_kernel(inputs_r, inputs_k, inputs_v, inputs_w, inputs_u)
+
+print(outputs_y.shape, outputs_y.sharding)
+print(y_state.shape, y_state.sharding)
+```
+
+---
+
+### rwkv6op 实现状态
+
+| Framework   | cuda | triton | native |
+|-------------|------|--------|--------|
+| PyTorch     | ✅   | ❌     | ✅     |
+| JAX         | ⚠️   | ❌     | ✅     |
+| TensorFlow  | ❌   | ❌     | ✅     |
+| NumPy       | ❌   | ❌     | ✅     |
+
+⚠️ JAX 的 CUDA 实现仅适用于 < 0.6.0，推荐 0.4.34。

{rwkv_ops-0.1.0.dist-info → rwkv_ops-0.2.dist-info}/RECORD

@@ -1,16 +1,20 @@
-rwkv_ops/__init__.py,sha256=Qp3EFYuuQB19ND6yvwxbczvGgH98hZcWr8_boIf_5I8,725
-rwkv_ops/rwkv7_kernel/__init__.py,sha256=nG_loc60Vp_ZQGLvwWFJYUIAsNCt-ldUYJnWdmVwyPc,5681
-rwkv_ops/rwkv7_kernel/get_jax_devices_info.py,sha256=fC_bISkSrKk1DvDNEaAdThmDZElz7T1BGWRl6EOXd6M,6005
-rwkv_ops/rwkv7_kernel/get_torch_devices_info.py,sha256=3Sgzw_-9VElBecmQ3TbW4y8GI9QZfwlP6BMb7CovlQ8,7380
-rwkv_ops/rwkv7_kernel/jax_op.py,sha256=xOg2YbPPHkXEJ59sOo1_stOrghJQPuj2nhXrqw4GSM4,9109
+rwkv_ops/__init__.py,sha256=Kfw_9iearbIplpuxx8sUl30TxKbTO-Ehqe0290Y4sZw,841
+rwkv_ops/rwkv6_kernel/__init__.py,sha256=_j6G_3fY8xPxrlZbgDT2ndX4IPiNJ4qjqIcdmNI_r9Q,4100
+rwkv_ops/rwkv6_kernel/jax_rwkv_kernel.py,sha256=WOzqfQQSHHMoWqm2kRz_BhtMzGYc5USJ26qaEwuARo4,30117
+rwkv_ops/rwkv6_kernel/ops_rwkv_kernel.py,sha256=otjfw5n6nf2YVpBIWIZjaCsxMyLXXwg-ma1ueXX-EdY,3274
+rwkv_ops/rwkv6_kernel/torch_rwkv_kernel.py,sha256=Q1uPMgaS21OEfQ8-sBDjaCUASMtkSOdN3OosEUsBp9U,12918
+rwkv_ops/rwkv7_kernel/__init__.py,sha256=GpwZ5dk7d5H6u-VSUAQ29KQVnCZdEwtKZ13_3kVJets,5888
+rwkv_ops/rwkv7_kernel/get_jax_devices_info.py,sha256=cMIaNED7d1PvYNSyq8wNI3G7wNvcgdUj9HWRBLuSVM8,6004
+rwkv_ops/rwkv7_kernel/get_torch_devices_info.py,sha256=ZL_rAM6lHB4nTOOU28Xm08qptfuIoijOMi_xwJG3KCo,7380
+rwkv_ops/rwkv7_kernel/jax_op.py,sha256=tyMxvk_EblDaGsePpxw3AhELvolp7LeE5NopUhKw1R0,9107
 rwkv_ops/rwkv7_kernel/native_keras_op.py,sha256=QPrXLbqw0chipQg_0jepRp2U19BYpBBFdKZWyaDNNoc,2488
-rwkv_ops/rwkv7_kernel/torch_op.py,sha256=elbv2SOHs9fzsO0FGxQ6ug7HVsjTZAE2oIM_ZOOY4JM,14451
+rwkv_ops/rwkv7_kernel/torch_op.py,sha256=yY5QP87iDow-T6a4ZzFShyIQ8gprTQoLYcjFqgOTW4Y,13675
 rwkv_ops/rwkv7_kernel/jax_kernel/__init__.py,sha256=uHsf_1qrtRK62IvhLuzefHGPWpHXmw1p0tqmwlHcptk,346
 rwkv_ops/rwkv7_kernel/jax_kernel/chunk_A_bwd.py,sha256=2Voq1Bdzn0DFloiLvwINBk7akmxRWIqXIQeyafrJJGg,2138
 rwkv_ops/rwkv7_kernel/jax_kernel/chunk_A_fwd.py,sha256=rhmglqHIIww7yPzaSBEp9ISxhhxoUbMtV51AUDyhUd8,1425
 rwkv_ops/rwkv7_kernel/jax_kernel/chunk_h_bwd.py,sha256=JDfVZsMb8yMlMN3sKT3i3l3y1YQiQkyUjnSNyan5Fqc,1888
 rwkv_ops/rwkv7_kernel/jax_kernel/chunk_h_fwd.py,sha256=g8b_81rIIjxeknYiklRGnox24rAvEvfKRKT-5nI0Euo,1992
-rwkv_ops/rwkv7_kernel/jax_kernel/chunk_o_bwd.py,sha256=W5gutV5WXA1bscfuDFCYNoRiG526CjHy4eaD2VtYuh0,3357
+rwkv_ops/rwkv7_kernel/jax_kernel/chunk_o_bwd.py,sha256=gQnToi1e1GZCvjWsEdWx6WakUN4Lc0JfaBSsSXYdN84,3369
 rwkv_ops/rwkv7_kernel/jax_kernel/chunk_o_fwd.py,sha256=4SjQ_zTZvFxsBMeWOx0JGFg9EQ4vllvEx30EcvSZJzI,853
 rwkv_ops/rwkv7_kernel/jax_kernel/cumsum.py,sha256=NoOh2_hA_rdH5bmaNNMAdCgVPfWvQpf-Q8BqF926jrw,667
 rwkv_ops/rwkv7_kernel/jax_kernel/wy_fast_bwd.py,sha256=PAMtE6wCW2Hz39oiHLGqhxY77csQAMYdNP2najDO_Jg,1407
@@ -20,7 +24,7 @@ rwkv_ops/rwkv7_kernel/torch_kernel/chunk_A_bwd.py,sha256=CWtotXkVvHz4-rkuOqWh6zK
 rwkv_ops/rwkv7_kernel/torch_kernel/chunk_A_fwd.py,sha256=4RJbyUTO23OxwH1rGVxeBiBVZKNHpPL_tJ7MFoDCIts,1475
 rwkv_ops/rwkv7_kernel/torch_kernel/chunk_h_bwd.py,sha256=zo6l0ZZUhXFu8wEFD76I0zSqFT9IXFKUKtyeaSwk380,1795
 rwkv_ops/rwkv7_kernel/torch_kernel/chunk_h_fwd.py,sha256=0ucN1U0EDTDqcyTPLLcsAX6FLTf2E_3toOY9p81gWYE,1858
-rwkv_ops/rwkv7_kernel/torch_kernel/chunk_o_bwd.py,sha256=CR6rK729Jfdt3JgZMpZxmALyTIYT6a7bmvllCm0GIxI,3281
+rwkv_ops/rwkv7_kernel/torch_kernel/chunk_o_bwd.py,sha256=ioPrS0NYQhpFk1j8rAxqtbwpx1CwjJQnrJEBDqVy-As,3283
 rwkv_ops/rwkv7_kernel/torch_kernel/chunk_o_fwd.py,sha256=54yoa3NpV64H-koURt-hUWpFHhUjwXpGvXPp2_ETCnw,825
 rwkv_ops/rwkv7_kernel/torch_kernel/cumsum.py,sha256=hQkpyaa0eUyB4V3UVks7l1_dHwOrbump0FZILityBKw,611
 rwkv_ops/rwkv7_kernel/torch_kernel/wy_fast_bwd.py,sha256=gk6QdoT1oq5B8Hp8Ak-SGqHm8CEj3MErUeWcRsaaOQM,1470
@@ -36,8 +40,8 @@ rwkv_ops/rwkv7_kernel/triton_kernel/cumsum.py,sha256=pRp_z587PrnpgRVpi031IndyjVI
 rwkv_ops/rwkv7_kernel/triton_kernel/utils.py,sha256=TNGlkwGq4t-TOcdVBk_N_vHPLzMFTu_F0V-O1RprIO4,553
 rwkv_ops/rwkv7_kernel/triton_kernel/wy_fast_bwd.py,sha256=szaG11q_WmpyhXi6aVWwzizvflCh5wND8wGA_V8afzA,5479
 rwkv_ops/rwkv7_kernel/triton_kernel/wy_fast_fwd.py,sha256=jbb19DUTHENU2RIOv_T4m_W1eXMqdRqG0XevIkBOhI4,9438
-rwkv_ops-0.1.0.dist-info/LICENSE.txt,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-rwkv_ops-0.1.0.dist-info/METADATA,sha256=KRIeem9Ulrk0bwVyelirLyWXcGVLoa1BXmwgjUPRikk,3572
-rwkv_ops-0.1.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-rwkv_ops-0.1.0.dist-info/top_level.txt,sha256=cVqoKE-WR_e2gHL87-6O4K1kG6-yTJGB2huyr6FmD2I,9
-rwkv_ops-0.1.0.dist-info/RECORD,,
+rwkv_ops-0.2.dist-info/LICENSE.txt,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+rwkv_ops-0.2.dist-info/METADATA,sha256=cUhC6EYLULLgNLVtLOg4qMSoAZjIDcrojax6egCTU04,8409
+rwkv_ops-0.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+rwkv_ops-0.2.dist-info/top_level.txt,sha256=cVqoKE-WR_e2gHL87-6O4K1kG6-yTJGB2huyr6FmD2I,9
+rwkv_ops-0.2.dist-info/RECORD,,

rwkv_ops-0.1.0.dist-info/METADATA

@@ -1,118 +0,0 @@
-Metadata-Version: 2.1
-Name: rwkv-ops
-Version: 0.1.0
-Home-page: https://github.com/your-org/rwkv_ops
-License: Apache 2.0
-Keywords: rwkv attention cuda triton pytorch jax
-Classifier: Development Status :: 3 - Alpha
-Classifier: Intended Audience :: Developers
-Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: Apache Software License
-Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Description-Content-Type: text/markdown
-License-File: LICENSE.txt
-Requires-Dist: keras
-
-[English Document](ENREADME.md)
-# RWKV OPS 项目
-> 由于 RWKV 将持续迭代，核心算子会随之更新。  
-> 本仓专门维护「算子」本身，不维护 layer 与 model；尽可能提供各框架的 GPU 算子。  
-> 目前：  
-> • GPU 算子：PyTorch、JAX（TensorFlow 待 Google 支持 Triton 后上线）  
-> • 原生算子：PyTorch、JAX、TensorFlow、NumPy  
-> 未来若 Keras 生态扩展，可能支持 MLX、OpenVINO。  
-> 注意：本库依赖 `keras`。
-
----
-
-## 环境变量
-
-| 变量名 | 含义 | 取值 | 默认值 | 优先级 |
-|---|---|---|---|---|
-| `KERAS_BACKEND` | Keras 后端 | jax / torch / tensorflow / numpy | — | 低 |
-| `KERNEL_BACKEND` | 算子后端 | jax / torch / tensorflow / numpy | torch | **高** |
-| `KERNEL_TYPE` | 实现类型 | triton / cuda / native | — | — |
-
-> 若 `KERNEL_BACKEND` 有值，直接采用；若为空，则用 `KERAS_BACKEND`；两者皆空则默认 torch。  
-> `native` 为原生算子，无 chunkwise，速度慢且显存高。
-
----
-
-## rwkv7op 使用方法
-
-```python
-from rwkv_ops import generalized_delta_rule  # 或 from rwkv_ops import rwkv7_op，完全等价
-
-def generalized_delta_rule(
-    r,
-    w,
-    k,
-    v,
-    a,
-    b,
-    initial_state=None,
-    output_final_state: bool = True,
-    head_first: bool = False,
-):
-    """
-    分块 Delta Rule 注意力接口。
-
-    Args:
-        q:  [B, T, H, K]
-        k:  [B, T, H, K]
-        v:  [B, T, H, V]
-        a:  [B, T, H, K]
-        b:  [B, T, H, K]
-        gk: [B, T, H, K]  # decay term in log space!
-        initial_state: 初始状态 [N, H, K, V]，N 为序列数
-        output_final_state: 是否返回最终状态
-        head_first: 是否 head-first 格式，不支持变长
-
-    Returns:
-        o:           输出 [B, T, H, V] 或 [B, H, T, V]
-        final_state: 最终状态 [N, H, K, V] 或 None
-    """
-```
-
----
-
-
-torch-cuda下head-size也是一个kernel参数，默认是64.
-若 head-size ≠ 64，请使用：
-
-```python
-from rwkv_ops import get_generalized_delta_rule
-
-generalized_delta_rule, RWKV7_USE_KERNEL = get_generalized_delta_rule(
-    your_head_size, KERNEL_TYPE="cuda"
-)
-```
-
-`RWKV7_USE_KERNEL` 为常量，标记是否使用 chunkwise 算子；
-因为两者padding 处理逻辑不同，具体如下
-
-```python
-if padding_mask is not None:
-    if RWKV7_USE_KERNEL:
-        w += (1 - padding_mask) * -1e9
-    else:
-        w = w * padding_mask + 1 - padding_mask
-```
-
----
-
-### rwkv7op的实现状态
-
-
-| Framework   | cuda | triton | native |
-|-------------|------|--------|--------|
-| PyTorch     | ✅   | ✅     | ✅     |
-| JAX         | ❌   | ✅     | ✅     | 
-| TensorFlow  | ❌   | ❌     | ✅     |
-| NumPy       | ❌   | ❌     | ✅     |