diff --git a/build/torch210-cu128-x86_64-windows/__init__.py b/build/torch210-cu128-x86_64-windows/__init__.py deleted file mode 100644 index cfd6a665271bbf1428c46acfb331ac4ff0258f5e..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch210-cu128-x86_64-windows/_ops.py b/build/torch210-cu128-x86_64-windows/_ops.py deleted file mode 100644 index d42b492624cadfc963e93ba90b8a1cdecc4b0999..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_cuda_38ccc47 -ops = torch.ops._rwkv_cuda_38ccc47 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_cuda_38ccc47::{op_name}" diff --git a/build/torch210-cu128-x86_64-windows/_rwkv_cuda_38ccc47.pyd b/build/torch210-cu128-x86_64-windows/_rwkv_cuda_38ccc47.pyd deleted file mode 100644 index 3ae6662c009f80685af977e90beb43c793b78c61..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/_rwkv_cuda_38ccc47.pyd +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e29ecb0b1d346a26decac4671633cf56fbc0b12df10027b1cca06a41f48976a7 -size 423424 diff --git a/build/torch210-cu128-x86_64-windows/metadata.json b/build/torch210-cu128-x86_64-windows/metadata.json deleted file mode 100644 index a1693688376cea25e6195d1985d05148d8e8d6fe..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/metadata.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "8.9", - "9.0" - ] - } -} diff --git a/build/torch210-cu128-x86_64-windows/rwkv/__init__.py b/build/torch210-cu128-x86_64-windows/rwkv/__init__.py deleted file mode 100644 index bc434ef44e63409acb52a8f3fff54a4adc46ed6a..0000000000000000000000000000000000000000 --- a/build/torch210-cu128-x86_64-windows/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch211-cu128-x86_64-windows/__init__.py b/build/torch211-cu128-x86_64-windows/__init__.py deleted file mode 100644 index cfd6a665271bbf1428c46acfb331ac4ff0258f5e..0000000000000000000000000000000000000000 --- a/build/torch211-cu128-x86_64-windows/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch211-cu128-x86_64-windows/_ops.py b/build/torch211-cu128-x86_64-windows/_ops.py deleted file mode 100644 index 93e271f508ad88a989f686c3b39bc87d9b750ee9..0000000000000000000000000000000000000000 --- a/build/torch211-cu128-x86_64-windows/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_cuda_34ead3e -ops = torch.ops._rwkv_cuda_34ead3e - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_cuda_34ead3e::{op_name}" diff --git a/build/torch211-cu128-x86_64-windows/_rwkv_cuda_34ead3e.pyd b/build/torch211-cu128-x86_64-windows/_rwkv_cuda_34ead3e.pyd deleted file mode 100644 index e83b455b115a355d558418433493a6ce6f1e48b2..0000000000000000000000000000000000000000 --- a/build/torch211-cu128-x86_64-windows/_rwkv_cuda_34ead3e.pyd +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:019c3000841ae541f82a5e49725fd28aafe20068e5b67685dd284da84d762148 -size 423424 diff --git a/build/torch211-cu128-x86_64-windows/metadata.json b/build/torch211-cu128-x86_64-windows/metadata.json deleted file mode 100644 index a1693688376cea25e6195d1985d05148d8e8d6fe..0000000000000000000000000000000000000000 --- a/build/torch211-cu128-x86_64-windows/metadata.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "8.9", - "9.0" - ] - } -} diff --git a/build/torch211-cu128-x86_64-windows/rwkv/__init__.py b/build/torch211-cu128-x86_64-windows/rwkv/__init__.py deleted file mode 100644 index d259f9535943d5617b4906d784ecd93a9d6dff82..0000000000000000000000000000000000000000 --- a/build/torch211-cu128-x86_64-windows/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch27-cxx11-cu118-x86_64-linux/rwkv/__init__.py b/build/torch27-cxx11-cu118-x86_64-linux/rwkv/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/rwkv/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch27-cxx11-cu118-x86_64-linux/rwkv/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-cu118-x86_64-linux/rwkv/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index c881de2e0f54ba27fce13d55f35013e2b90dea0b..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu118-x86_64-linux/rwkv/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu118-x86_64-linux/rwkv/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-cu118-x86_64-linux/rwkv/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index 49bded237b33ffeb440b398ca5da96497d92364c..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu118-x86_64-linux/rwkv/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu118-x86_64-linux/rwkv/_ops.py b/build/torch27-cxx11-cu118-x86_64-linux/rwkv/_ops.py deleted file mode 100644 index c54b92a73953ef7168a8a45b1d3caa2e64acc4cb..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/rwkv/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_0aa7cb0 -ops = torch.ops._rwkv_0aa7cb0 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_0aa7cb0::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu118-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so b/build/torch27-cxx11-cu118-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so deleted file mode 100644 index af67a351934b49c158ae3e5d7a4fa465ace19ab1..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a561edd91db3da01aad4f484eca06f44e98843f859317aeaa9a3df85339ce850 -size 2065400 diff --git a/build/torch27-cxx11-cu126-x86_64-linux/rwkv/__init__.py b/build/torch27-cxx11-cu126-x86_64-linux/rwkv/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/rwkv/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch27-cxx11-cu126-x86_64-linux/rwkv/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-cu126-x86_64-linux/rwkv/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index 7dee5151cb6937785bd5a63c6a8f02f9bb3d351c..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu126-x86_64-linux/rwkv/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu126-x86_64-linux/rwkv/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-cu126-x86_64-linux/rwkv/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index 3bad2999c65772f5fd868bf6af0fc88b592fa665..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu126-x86_64-linux/rwkv/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu126-x86_64-linux/rwkv/_ops.py b/build/torch27-cxx11-cu126-x86_64-linux/rwkv/_ops.py deleted file mode 100644 index c54b92a73953ef7168a8a45b1d3caa2e64acc4cb..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/rwkv/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_0aa7cb0 -ops = torch.ops._rwkv_0aa7cb0 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_0aa7cb0::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu126-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so b/build/torch27-cxx11-cu126-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so deleted file mode 100644 index fdf4405e0825430c46f683f9a763da3a39d96873..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b115095e24c64d21eecefdc0ed4d9c13ea9485930f8f56cc99874762c501f278 -size 2106408 diff --git a/build/torch27-cxx11-cu128-x86_64-linux/rwkv/__init__.py b/build/torch27-cxx11-cu128-x86_64-linux/rwkv/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/rwkv/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch27-cxx11-cu128-x86_64-linux/rwkv/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-cu128-x86_64-linux/rwkv/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index 382a6ba0237ee4e4fe73cbb610a447987d9005e7..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu128-x86_64-linux/rwkv/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu128-x86_64-linux/rwkv/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-cu128-x86_64-linux/rwkv/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index 0fdac79c69dd79159f35740b56bde8c18f82a375..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu128-x86_64-linux/rwkv/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu128-x86_64-linux/rwkv/_ops.py b/build/torch27-cxx11-cu128-x86_64-linux/rwkv/_ops.py deleted file mode 100644 index c54b92a73953ef7168a8a45b1d3caa2e64acc4cb..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/rwkv/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_0aa7cb0 -ops = torch.ops._rwkv_0aa7cb0 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_0aa7cb0::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu128-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so b/build/torch27-cxx11-cu128-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so deleted file mode 100644 index df248c48b997c805b9ec6368141f1d027a9c641b..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/rwkv/_rwkv_0aa7cb0.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:19029bb516a8cf3428f4e47a2b5f26f091d169206e51df7c4e89880553a15df7 -size 2308848 diff --git a/build/torch28-cxx11-cu126-x86_64-linux/__init__.py b/build/torch28-cxx11-cu126-x86_64-linux/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch28-cxx11-cu126-x86_64-linux/_ops.py b/build/torch28-cxx11-cu126-x86_64-linux/_ops.py deleted file mode 100644 index ada2ab024f61c605a79ddc243163cfb3b8a39ef7..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_44f2fa4 -ops = torch.ops._rwkv_44f2fa4 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_44f2fa4::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu126-x86_64-linux/_rwkv_44f2fa4.abi3.so b/build/torch28-cxx11-cu126-x86_64-linux/_rwkv_44f2fa4.abi3.so deleted file mode 100644 index 4cc06763b7d171a586c906cf1e6bb55cba0e4709..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/_rwkv_44f2fa4.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0d383c1186dde17bb427291205ae83bdb2e49abf15b4eb482151aa0940e51d4d -size 2111504 diff --git a/build/torch28-cxx11-cu126-x86_64-linux/metadata.json b/build/torch28-cxx11-cu126-x86_64-linux/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-cu126-x86_64-linux/rwkv/__init__.py b/build/torch28-cxx11-cu126-x86_64-linux/rwkv/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch28-cxx11-cu128-x86_64-linux/__init__.py b/build/torch28-cxx11-cu128-x86_64-linux/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch28-cxx11-cu128-x86_64-linux/_ops.py b/build/torch28-cxx11-cu128-x86_64-linux/_ops.py deleted file mode 100644 index ada2ab024f61c605a79ddc243163cfb3b8a39ef7..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_44f2fa4 -ops = torch.ops._rwkv_44f2fa4 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_44f2fa4::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu128-x86_64-linux/_rwkv_44f2fa4.abi3.so b/build/torch28-cxx11-cu128-x86_64-linux/_rwkv_44f2fa4.abi3.so deleted file mode 100644 index 026e5fa245ea37eee14a4233f948fd28abe48747..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/_rwkv_44f2fa4.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5fb0335b36148895aac63e1822e14c4be054dd842dbdddcfab668b65b2335ff2 -size 2318032 diff --git a/build/torch28-cxx11-cu128-x86_64-linux/metadata.json b/build/torch28-cxx11-cu128-x86_64-linux/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-cu128-x86_64-linux/rwkv/__init__.py b/build/torch28-cxx11-cu128-x86_64-linux/rwkv/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch28-cxx11-cu129-x86_64-linux/__init__.py b/build/torch28-cxx11-cu129-x86_64-linux/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch28-cxx11-cu129-x86_64-linux/_ops.py b/build/torch28-cxx11-cu129-x86_64-linux/_ops.py deleted file mode 100644 index ada2ab024f61c605a79ddc243163cfb3b8a39ef7..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_44f2fa4 -ops = torch.ops._rwkv_44f2fa4 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_44f2fa4::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu129-x86_64-linux/_rwkv_44f2fa4.abi3.so b/build/torch28-cxx11-cu129-x86_64-linux/_rwkv_44f2fa4.abi3.so deleted file mode 100644 index 62492882cb398511bfca1748b134d07316503937..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/_rwkv_44f2fa4.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2178eb41ec85f3861d61391de45ec47c3127490fbd902f1e2167683514cf867b -size 2335432 diff --git a/build/torch28-cxx11-cu129-x86_64-linux/metadata.json b/build/torch28-cxx11-cu129-x86_64-linux/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-cu129-x86_64-linux/rwkv/__init__.py b/build/torch28-cxx11-cu129-x86_64-linux/rwkv/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/__init__.py b/build/torch29-cxx11-cu126-aarch64-linux/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-aarch64-linux/_ops.py b/build/torch29-cxx11-cu126-aarch64-linux/_ops.py deleted file mode 100644 index 5dffb89844d9c7644252ba9bc5b0635b106174cd..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_cuda_efd954c -ops = torch.ops._rwkv_cuda_efd954c - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_cuda_efd954c::{op_name}" diff --git a/build/torch29-cxx11-cu126-aarch64-linux/_rwkv_cuda_efd954c.abi3.so b/build/torch29-cxx11-cu126-aarch64-linux/_rwkv_cuda_efd954c.abi3.so deleted file mode 100644 index 894d88388663606729b454e35fd4a17f9de139dc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/_rwkv_cuda_efd954c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f6d39263049bd9e5823007907a9b2f366226c9dacc5a6157f68f8e317bd1108c -size 2230968 diff --git a/build/torch29-cxx11-cu126-aarch64-linux/metadata.json b/build/torch29-cxx11-cu126-aarch64-linux/metadata.json deleted file mode 100644 index 724a12daaf58a25e6b5fc3e62fc3b76e4b54ccd9..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/metadata.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "8.0", - "8.9", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu126-aarch64-linux/rwkv/__init__.py b/build/torch29-cxx11-cu126-aarch64-linux/rwkv/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu126-x86_64-linux/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_ops.py b/build/torch29-cxx11-cu126-x86_64-linux/_ops.py deleted file mode 100644 index ada2ab024f61c605a79ddc243163cfb3b8a39ef7..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_44f2fa4 -ops = torch.ops._rwkv_44f2fa4 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_44f2fa4::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_rwkv_44f2fa4.abi3.so b/build/torch29-cxx11-cu126-x86_64-linux/_rwkv_44f2fa4.abi3.so deleted file mode 100644 index 87947ce206483ed63225e5ac32379232817c05b3..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/_rwkv_44f2fa4.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:64e542cd9e7f9a7ea5bf22a91a27fb34f6352b7fbbc39463f2f5d183f6386d05 -size 2111480 diff --git a/build/torch29-cxx11-cu126-x86_64-linux/metadata.json b/build/torch29-cxx11-cu126-x86_64-linux/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/rwkv/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/rwkv/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-aarch64-linux/_ops.py b/build/torch29-cxx11-cu128-aarch64-linux/_ops.py deleted file mode 100644 index 5dffb89844d9c7644252ba9bc5b0635b106174cd..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_cuda_efd954c -ops = torch.ops._rwkv_cuda_efd954c - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_cuda_efd954c::{op_name}" diff --git a/build/torch29-cxx11-cu128-aarch64-linux/_rwkv_cuda_efd954c.abi3.so b/build/torch29-cxx11-cu128-aarch64-linux/_rwkv_cuda_efd954c.abi3.so deleted file mode 100644 index 87d47bd77178df8cfa9089434aa9eca8f59562de..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/_rwkv_cuda_efd954c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d1cde967371ab4341efc9a43eead44a82fe2e80099f170e0f4781f96f8d68961 -size 2427688 diff --git a/build/torch29-cxx11-cu128-aarch64-linux/metadata.json b/build/torch29-cxx11-cu128-aarch64-linux/metadata.json deleted file mode 100644 index 890b88c81c1da1a10714a26cdca182aa27c63f6f..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/metadata.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "8.9", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu128-aarch64-linux/rwkv/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/rwkv/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py b/build/torch29-cxx11-cu128-x86_64-linux/_ops.py deleted file mode 100644 index ada2ab024f61c605a79ddc243163cfb3b8a39ef7..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_44f2fa4 -ops = torch.ops._rwkv_44f2fa4 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_44f2fa4::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_rwkv_44f2fa4.abi3.so b/build/torch29-cxx11-cu128-x86_64-linux/_rwkv_44f2fa4.abi3.so deleted file mode 100644 index 602898b1e3cb5930b41c1772b31a77942b830400..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/_rwkv_44f2fa4.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e37b16a151528c1ad0eecfeb1c3296ac45e6856d3b5b08f7acc68499e3ceca4c -size 2318000 diff --git a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json b/build/torch29-cxx11-cu128-x86_64-linux/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-x86_64-linux/rwkv/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/rwkv/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/__init__.py b/build/torch29-cxx11-cu129-aarch64-linux/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch29-cxx11-cu129-aarch64-linux/_ops.py b/build/torch29-cxx11-cu129-aarch64-linux/_ops.py deleted file mode 100644 index aab83e9761559fc94019184b65cc97ea30c6b43e..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_cuda_3798a3c -ops = torch.ops._rwkv_cuda_3798a3c - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_cuda_3798a3c::{op_name}" diff --git a/build/torch29-cxx11-cu129-aarch64-linux/_rwkv_cuda_3798a3c.abi3.so b/build/torch29-cxx11-cu129-aarch64-linux/_rwkv_cuda_3798a3c.abi3.so deleted file mode 100644 index de14609ba26012446d2473ae5066ed554999df07..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/_rwkv_cuda_3798a3c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:cd7fa25260f4e9d6506b1c9c56009144f20dc08d042f832f87312dbd78d62f7c -size 2429112 diff --git a/build/torch29-cxx11-cu129-aarch64-linux/metadata.json b/build/torch29-cxx11-cu129-aarch64-linux/metadata.json deleted file mode 100644 index 890b88c81c1da1a10714a26cdca182aa27c63f6f..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/metadata.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "8.9", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu129-aarch64-linux/rwkv/__init__.py b/build/torch29-cxx11-cu129-aarch64-linux/rwkv/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch29-cxx11-cu129-x86_64-linux/_ops.py b/build/torch29-cxx11-cu129-x86_64-linux/_ops.py deleted file mode 100644 index aab83e9761559fc94019184b65cc97ea30c6b43e..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_cuda_3798a3c -ops = torch.ops._rwkv_cuda_3798a3c - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_cuda_3798a3c::{op_name}" diff --git a/build/torch29-cxx11-cu129-x86_64-linux/_rwkv_cuda_3798a3c.abi3.so b/build/torch29-cxx11-cu129-x86_64-linux/_rwkv_cuda_3798a3c.abi3.so deleted file mode 100644 index 282967d545b3cf1a383cd1f00e18228f8e9a8e97..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/_rwkv_cuda_3798a3c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2dc16b9b2d7deddc8a2cc6747cab332c9950ea3404c946446165ee159d322267 -size 2335432 diff --git a/build/torch29-cxx11-cu129-x86_64-linux/metadata.json b/build/torch29-cxx11-cu129-x86_64-linux/metadata.json deleted file mode 100644 index 890b88c81c1da1a10714a26cdca182aa27c63f6f..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/metadata.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "8.9", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu129-x86_64-linux/rwkv/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/rwkv/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch29-cxx11-cu130-aarch64-linux/_ops.py b/build/torch29-cxx11-cu130-aarch64-linux/_ops.py deleted file mode 100644 index 5dffb89844d9c7644252ba9bc5b0635b106174cd..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_cuda_efd954c -ops = torch.ops._rwkv_cuda_efd954c - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_cuda_efd954c::{op_name}" diff --git a/build/torch29-cxx11-cu130-aarch64-linux/_rwkv_cuda_efd954c.abi3.so b/build/torch29-cxx11-cu130-aarch64-linux/_rwkv_cuda_efd954c.abi3.so deleted file mode 100644 index 8325463d49dec8d7492c677dbe0863e4056c206f..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/_rwkv_cuda_efd954c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ab7e3bbeaf2bb811191f7dd30fa68f7c5c772423020b9fe994d31273bf77dc4a -size 2431672 diff --git a/build/torch29-cxx11-cu130-aarch64-linux/metadata.json b/build/torch29-cxx11-cu130-aarch64-linux/metadata.json deleted file mode 100644 index 890b88c81c1da1a10714a26cdca182aa27c63f6f..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/metadata.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "12.0", - "8.0", - "8.9", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu130-aarch64-linux/rwkv/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/rwkv/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/__init__.py deleted file mode 100644 index d4ea22dbc160d2b75b32c3c8c5d882d25c77b494..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/__init__.py +++ /dev/null @@ -1,170 +0,0 @@ -from ._ops import ops -from typing import Tuple, Any - -# Use a broad Tensor alias to avoid importing torch at import time. -from torch import Tensor - -def forward(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (float32). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward(w, u, k, v, y) - - -def forward_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor) -> None: - """RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``). - - Runs the CUDA kernel and writes the result into ``y`` in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y)) - ops.forward_bf16(w, u, k, v, y) - - -def forward_with_state(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (float32). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state(w, u, k, v, y, s) - - -def forward_with_state_bf16(w: Tensor, u: Tensor, k: Tensor, v: Tensor, y: Tensor, s: Tensor) -> None: - """RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``). - - Runs the CUDA kernel using and updating state ``s`` and writes the result into ``y``. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs. - """ - _validate_device_match((w, u, k, v, y, s)) - ops.forward_with_state_bf16(w, u, k, v, y, s) - - -def backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (float32). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def backward_bf16( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - y: Tensor, - gy: Tensor, - gw: Tensor, - gu: Tensor, - gk: Tensor, - gv: Tensor, -) -> None: - """RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``). - - Writes gradients into the provided tensors in-place. - - Args: - w: Decay weights, shape ``[C]``, dtype ``torch.float32``. - u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``. - gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place). - gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place). - - Notes: - - All tensors must be on the same CUDA device. - - Shapes must agree on ``B``, ``T`` and ``C``. - """ - _validate_device_match((w, u, k, v, y, gy, gw, gu, gk, gv)) - ops.backward_bf16(w, u, k, v, y, gy, gw, gu, gk, gv) - - -def _validate_device_match(tensors: Tuple[Tensor, ...]) -> None: - """Minimal runtime validation that all tensors live on the same CUDA device.""" - if not tensors: - return - device = tensors[0].device - if not device.type == "cuda": - raise RuntimeError("RWKV CUDA ops require CUDA tensors") - for t in tensors[1:]: - if t.device != device: - raise RuntimeError("All tensors must be on the same CUDA device") - - -__all__ = [ - "forward", - "forward_bf16", - "forward_with_state", - "forward_with_state_bf16", - "backward", - "backward_bf16", -] \ No newline at end of file diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py b/build/torch29-cxx11-cu130-x86_64-linux/_ops.py deleted file mode 100644 index ada2ab024f61c605a79ddc243163cfb3b8a39ef7..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _rwkv_44f2fa4 -ops = torch.ops._rwkv_44f2fa4 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_rwkv_44f2fa4::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_rwkv_44f2fa4.abi3.so b/build/torch29-cxx11-cu130-x86_64-linux/_rwkv_44f2fa4.abi3.so deleted file mode 100644 index 6844b7973cff37a87fa06f89130dba801198b046..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/_rwkv_44f2fa4.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:66177f2e067a2624a79e80664f3d87547c1a5916a91711f792495c4a77f5b70f -size 2343392 diff --git a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json b/build/torch29-cxx11-cu130-x86_64-linux/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch29-cxx11-cu130-x86_64-linux/rwkv/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/rwkv/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/rwkv/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))