diff --git a/.gitattributes b/.gitattributes index 715ca49acf9f24c9a7968c932373b3dd4663f662..da3f0c04fb9096fcd38e3e313272f98982136af4 100644 --- a/.gitattributes +++ b/.gitattributes @@ -252,3 +252,21 @@ build/torch213-cxx11-cu132-x86_64-linux/_rwkv_cuda_89d4054.abi3.so filter=lfs di build/torch214-cxx11-cu126-x86_64-linux/_rwkv_cuda_89d4054.abi3.so filter=lfs diff=lfs merge=lfs -text build/torch214-cxx11-cu130-x86_64-linux/_rwkv_cuda_89d4054.abi3.so filter=lfs diff=lfs merge=lfs -text build/torch214-cxx11-cu132-x86_64-linux/_rwkv_cuda_89d4054.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch210-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch210-cu128-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch210-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch211-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch211-cu128-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch211-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch212-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch212-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch212-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch213-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch213-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch213-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch214-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch214-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch214-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch215-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch215-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch215-cu134-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text diff --git a/build/torch210-cu126-aarch64-linux/__init__.py b/build/torch210-cu126-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch210-cu126-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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-cu126-aarch64-linux/_ops.py b/build/torch210-cu126-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch210-cu126-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch210-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch210-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..a27c59da62fc7e67ef28a2e7412175f929c76db7 --- /dev/null +++ b/build/torch210-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6c47376791be5ee676c27115135281c658da006ef720dc42d8655a02681db8be +size 286296 diff --git a/build/torch210-cu126-aarch64-linux/metadata.json b/build/torch210-cu126-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..6340eb135576410388307203a035fbc2ef4df7d5 --- /dev/null +++ b/build/torch210-cu126-aarch64-linux/metadata.json @@ -0,0 +1,37 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "bEc3Z5G+XuZ2wnEVE1KBxljaAG73INxC2GVaAmgduL4=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch210-cu126-aarch64-linux/metadata.json.sigstore b/build/torch210-cu126-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..e27a481c467b4e525ef21e993a3793f204bdad51 --- /dev/null +++ b/build/torch210-cu126-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHKzCCBrGgAwIBAgIUZSVhv0na+qJH0Nf6L/XxFjBMIuAwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTEzWhcNMjYxMDE5MTMxNTEzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEIEXwVoqUK2TFzfX5g8hzcwXxK4S+x3htZE7eJxD5Vwz+9YhFKGERil972rFDeXwyT/a6iGjPPLg3kutV0pi1faOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVEQiOgAABAMARjBEAiAtFwMgyO6REwpAjWzE+eTZGcPQL1/JtRG+LVz/LrEpIgIgbcACB4jfoVa0mdHWGohHx0rY3ecys67hO29q53XFA6QwCgYIKoZIzj0EAwMDaAAwZQIxAJMfz0KAtHElbfGLdxDwQPu3kf2kJ/FzZKYtkCCMnemyeiNtmSJXaD+7LpUIaqANiwIwALOLOMWnFnJqTZqreLLFq2MkN84yxC27Xq6t3jJHMRYkMjG8T0x7XuDP9jPnuL9B"},"tlogEntries":[{"logIndex":"3321255773","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415114","inclusionPromise":{"signedEntryTimestamp":"MEUCIB5qXxvEP6a3rFbi5kT4ENatp6YGfV+mBlRd0mbSk4V2AiEA0Otb65MLSEe5m2ztjcl7h+JSqB2AERLjiLATRGATirc="},"inclusionProof":{"logIndex":"3199351511","rootHash":"aPii2mUViYuXwH4NxQ5BqtWdykR9rabadRfKzrMogbU=","treeSize":"3199351519","hashes":["R5yA4Qum/KSkAjsH0FlyzJBTut0Qo8eNCr+ET6CuacA=","/5vL3V1p+ras6WV27ZUUdowXCJjXHs0s1PABdsWZNDY=","cysFCrOSao4HP51eNi3jrIzgyF/MHu3DxbjlUFP3pGk=","PTRxdspNiGQy87LfUWpwE4t3C+TmZPFKN5HS/tMq5Zc=","czgfccxT5XQkQATWj1yhgeWemmgBo0qm8m+0fWQcCmE=","C2wlB1njZGHSG8+4TOxjL1AyDGU84s1en2l17JA/gio=","cKJ2x4/0Tl3LMrC4gHJ8RTMxrNl2gYom9/44N8hNOI0=","pkUNmvOtKnFBM2KMw42O+V67fz6YzLuD4OSqzBsCopI=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199351519\naPii2mUViYuXwH4NxQ5BqtWdykR9rabadRfKzrMogbU=\n\n— rekor.sigstore.dev wNI9ajBGAiEA7KDEykUW5zc7TkSMhRYIAC+a5Rz7LUZhj9g3f6IkkAECIQCchoH0QEZ0BIzVXwzlcnJKPLVRIrU1cBmQOJDXHeApUQ==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIxM2M3NjVhOGEyZmZmODg5MjgxNzI3MDQwNGI1YjU5ZWNkNmIzYzU0ZWMwODNhYjQzYWI5MDdjMDIzNzkwNTBjIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJQTJCTFJoODlKNEZ6ZkVWZ1ZRMjBvUFRsTFJ5MEdKaC9XWHdya3BTbFUyQkFpQStrYnJTRDUxdGhGb1RIZTZ2L2ZjdkJoR2ZQL2ZrLyszempQZGlRK3RoZ2c9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExla05EUW5KSFowRjNTVUpCWjBsVldsTldhSFl3Ym1FcmNVcElNRTVtTmt3dldIaEdha0pOU1hWQmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJGZWxkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVWNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZKUlZoM1ZtOXhWVXN5VkVaNlpsZzFaemhvZW1OM1dIaExORk1yZUROb2RGcEZOMlVLU25oRU5WWjNlaXM1V1doR1MwZEZVbWxzT1RjeWNrWkVaVmgzZVZRdllUWnBSMnBRVUV4bk0ydDFkRll3Y0dreFptRlBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWUmFVOW5RVUZDUVUxQlVtcENSVUZwUVhSR2QwMW5lVTgyVWtWM2NFRnFWM3BGSzJWVVdrZGpVRkZNTVM5S2RGSkhLd3BNVm5vdlRISkZjRWxuU1dkaVkwRkRRalJxWm05V1lUQnRaRWhYUjI5b1NIZ3djbGt6WldONWN6WTNhRTh5T1hFMU0xaEdRVFpSZDBObldVbExiMXBKQ25wcU1FVkJkMDFFWVVGQmQxcFJTWGhCU2sxbWVqQkxRWFJJUld4aVprZE1aSGhFZDFGUWRUTnJaakpyU2k5R2VscExXWFJyUTBOTmJtVnRlV1ZwVG5RS2JWTktXR0ZFS3pkTWNGVkpZWEZCVG1sM1NYZEJURTlNVDAxWGJrWnVTbkZVV25GeVpVeE1SbkV5VFd0T09EUjVlRU15TjFoeE5uUXpha3BJVFZKWmF3cE5ha2M0VkRCNE4xaDFSRkE1YWxCdWRVdzVRZ290TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyDADAgEAMIICvwYJKoZIhvcNAQcCoIICsDCCAqwCAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgBqHOvrVBjkRjHaFAi5yt98nxIFFHyuZCTdHzR3stX+4CFEZgy8jMHsbSAGyW/ekF66Sh9t7kGA8yMDI2MTAxOTEzMDUxM1owAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdowggHWAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTMwNTEzWjAvBgkqhkiG9w0BCQQxIgQgGBBWrK32S3QdKxdD9Hh7TTH1ZypiDDj7vkcPviIztAIwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGYwZAIwVRm3NhFlA7D5l20zlhqMV8DcmRB26YSIXCk/n/VSSUCUaRxDlB5HohpPY9JzA7xgAjBH8e3d0AYSnNbQWeXIoGLVLwYEb4Prp1cI2Eyexall8sKca11GQovIzYrUp6du6+Q="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"E8dlqKL/+IkoFycEBLW1ns1rPFTsCDq0OrkHwCN5BQw="},"signature":"MEQCIA2BLRh89J4FzfEVgVQ20oPTlLRy0GJh/WXwrkpSlU2BAiA+kbrSD51thFoTHe6v/fcvBhGfP/fk/+3zjPdiQ+thgg=="}} \ No newline at end of file diff --git a/build/torch210-cu126-aarch64-linux/symbols.json b/build/torch210-cu126-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch210-cu126-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch210-cu128-aarch64-linux/__init__.py b/build/torch210-cu128-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch210-cu128-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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-aarch64-linux/_ops.py b/build/torch210-cu128-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch210-cu128-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch210-cu128-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch210-cu128-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..868a3f46f4d824ac95177561747121c3725ef2f5 --- /dev/null +++ b/build/torch210-cu128-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:14be011dc6e8ccec33a280a3bc9feeca7491314ad957b18b77261ff4c3fa1538 +size 417376 diff --git a/build/torch210-cu128-aarch64-linux/metadata.json b/build/torch210-cu128-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..22ced57f6649ae1a5098c68fa23e1fc8799a51a9 --- /dev/null +++ b/build/torch210-cu128-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "FL4BHcbozOwzooCjvJ/uynSRMUrZV7GLdyYf9MP6FTg=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch210-cu128-aarch64-linux/metadata.json.sigstore b/build/torch210-cu128-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..5e5ed0f0d335b439385c8ab17adc3897426db1ae --- /dev/null +++ b/build/torch210-cu128-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrOgAwIBAgIUZ6HEQ502b3n2puuivfEs3VvqblUwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTE0WhcNMjYxMDE5MTMxNTE0WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEvLGuTOfyYvjQ1dN2N4wbKyXSSvvK4ets18+ykHB4PNnKzCejz8SkdC4GQhBLnPBiG2PkBvGpKRzTc6M/jpcKPKOCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVEQmDwAABAMASDBGAiEAlGHTrNqKDVGPzXJt4ymjVqFWmpn3QM5tYha+x+nQ6OoCIQDHkp6/vLmN96oesoohOMG14sihg4qeR6F9SyE04VTydjAKBggqhkjOPQQDAwNoADBlAjA7lvNAhLrHRUF1J6x5yar/yOTfYFN+RXye7DbRV59JUwDWEr/Rks6VCYgpUYQFUDkCMQDTzObuMtzrEcMTiiesKt1Tg/hX01pu70xgX6QloY2k1/3n4ahCotbaBbFC5D2LePQ="},"tlogEntries":[{"logIndex":"3321255814","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415114","inclusionPromise":{"signedEntryTimestamp":"MEUCIG3l5uGRjkKthUEJoQ49ibP9WDOFqpW03WjNha8wY6m8AiEA3ZPu+0pImIXK0fBSB82k8N0KLNYXV7Uk9FrAcH/pesE="},"inclusionProof":{"logIndex":"3199351552","rootHash":"QI94EYycUBDhW6Y+CDbIPOEX6B2B8f+qzFCF+5HzmhM=","treeSize":"3199351562","hashes":["7ykgO/88CqExh15GgQYoAtkYan3BK2rMOj0gFPu6QXE=","68HM1DQoi4LC3i+K/6Vrggq4+bKDe7w+OGQ8hKU8ePU=","wFL2EATaHCpKyDC9FBHeDZ4sG/0LQQfCtFuxTJ/TnAw=","QGoej8LxZdrPR9krEM4XUeok+BBSNe7f3aElK/vBCpk=","xqRGIj7vKKnGXGhCxIxXy1I2FKxqmR2s2WTICdDM2+8=","pkUNmvOtKnFBM2KMw42O+V67fz6YzLuD4OSqzBsCopI=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199351562\nQI94EYycUBDhW6Y+CDbIPOEX6B2B8f+qzFCF+5HzmhM=\n\n— rekor.sigstore.dev wNI9ajBEAiA/h8i8ezEjHP2Ls+1KuUK1g10bgXw1Ou6f/tksEMM3JQIgXW/h70zKk8fcwxfd3JblyVsPwExT0Yp+pm2VQHHJ87g=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJiOGNlZWNiYWRmYjhhMTJhNGY2ZTBlMWFlMjBhNWUxNDQyYjYwMzU2YzEwZmJjNDQyYzgxNGE0ZjIxMGFkMDY4In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUURLT0w3Y1ZZaDF1RXdBTWN2dmw3YnVzVlJRVy96bGJncmZ5NGc0eFBoamtBSWhBSXFwUFYxMTNWTE81b2Nmbllkdi9IUmcwWHMwWWlZdUVKdktWZGpHM3FzUyIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KUFowRjNTVUpCWjBsVldqWklSVkUxTURKaU0yNHljSFYxYVhabVJYTXpWblp4WW14VmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJGTUZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVVd1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVYyVEVkMVZFOW1lVmwyYWxFeFpFNHlUalIzWWt0NVdGTlRkblpMTkdWMGN6RTRLM2tLYTBoQ05GQk9ia3Q2UTJWcWVqaFRhMlJETkVkUmFFSk1ibEJDYVVjeVVHdENka2R3UzFKNlZHTTJUUzlxY0dOTFVFdFBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWUmJVUjNRVUZDUVUxQlUwUkNSMEZwUlVGc1IwaFVjazV4UzBSV1IxQjZXRXAwTkhsdGFsWnhSbGR0Y0c0elVVMDFkQXBaYUdFcmVDdHVVVFpQYjBOSlVVUklhM0EyTDNaTWJVNDVObTlsYzI5dmFFOU5SekUwYzJsb1p6UnhaVkkyUmpsVGVVVXdORlpVZVdScVFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2IwRkVRbXhCYWtFM2JIWk9RV2hNY2toU1ZVWXhTalo0TlhsaGNpOTVUMVJtV1VaT0sxSlllV1UzUkdKU1ZqVTVTbFYzUkZjS1JYSXZVbXR6TmxaRFdXZHdWVmxSUmxWRWEwTk5VVVJVZWs5aWRVMTBlbkpGWTAxVWFXbGxjMHQwTVZSbkwyaFlNREZ3ZFRjd2VHZFlObEZzYjFreWF3b3hMek51TkdGb1EyOTBZbUZDWWtaRE5VUXlUR1ZRVVQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgUytML7vytMJU9imxLRpfrSmEkWrAm+TgUfZnE3iphxgCFQDmcq7Mu0cNRzEPMc/RI9ZGNehegRgPMjAyNjEwMTkxMzA1MTRaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTEzMDUxNFowLwYJKoZIhvcNAQkEMSIEIPKZSagTbrkepqFJkFasUHPDo+N8rtr56qR7w4QHHor7MIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMBDY0WsGwZakLzAhn17wc29LK3hkUpON84RdgkzcvivNlhUWWNZl8I+p/bYtXj10/gIwGGdcRpP69s1XJ0WkJ/jfEs3J3z7MiY5t1yBl0Hceb82jXmYL07Iojiuasw4jgMpC"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"uM7sut+4oSpPbg4a4gpeFEK2A1bBD7xELIFKTyEK0Gg="},"signature":"MEYCIQDKOL7cVYh1uEwAMcvvl7busVRQW/zlbgrfy4g4xPhjkAIhAIqpPV113VLO5ocfnYdv/HRg0Xs0YiYuEJvKVdjG3qsS"}} \ No newline at end of file diff --git a/build/torch210-cu128-aarch64-linux/symbols.json b/build/torch210-cu128-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch210-cu128-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch210-cu130-aarch64-linux/__init__.py b/build/torch210-cu130-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch210-cu130-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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-cu130-aarch64-linux/_ops.py b/build/torch210-cu130-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch210-cu130-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch210-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch210-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..d136ea77af43c88fc25c75d30838a06b08c0fb0a --- /dev/null +++ b/build/torch210-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bfffd3c0c5a267765d1216cdbecaed258220103af006537aba08dd419f9c45bc +size 421472 diff --git a/build/torch210-cu130-aarch64-linux/metadata.json b/build/torch210-cu130-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..71b9fc178fe449a8ecdbf6b9fc5ac8e4775cef9a --- /dev/null +++ b/build/torch210-cu130-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "v//TwMWiZ3ZdEhbNvsrtJYIgEDrwBlN6ugjdQZ+cRbw=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch210-cu130-aarch64-linux/metadata.json.sigstore b/build/torch210-cu130-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..78227ad3c3fa2caead997dc43a4cef6d61513225 --- /dev/null +++ b/build/torch210-cu130-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrKgAwIBAgIUb0LpjxrK0Maa+ZODgUg6KoRnZCIwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTE1WhcNMjYxMDE5MTMxNTE1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE8sjer4q7SDt3umHi5iqVbVgvlPqgDYqB7sLo4+zQwNw8agBR3wNgirNoDZcSwFtP53inVwhUanEIaGRcodYHaqOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVEQpegAABAMARzBFAiEAtenKv2UoTuOdPCjot60HUABQ9VGzVg0jT/NoFB1mM0YCID+2z08Ud5yeVbj8ipNNu034eV9yTiKrwvMTYDH0zRXBMAoGCCqGSM49BAMDA2gAMGUCMDDyAEBce8fPqwGs4TRuYQVK/lb+eGCZOu4LvOFi8BsiKIEkRMYQE1BxJx7BEm+wRQIxALHR31InaWblMmSZuNoTatMG9TBWcQAjX4gf8EFqEc3vfbqBNb5f1AOkllapQbX1dg=="},"tlogEntries":[{"logIndex":"3321255861","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415115","inclusionPromise":{"signedEntryTimestamp":"MEQCIChCoQdSZLNcwv7ZPdmBF0Cq7nInYhSTLETsYp/YlqMXAiA16Z/9EoaSmg51Yl3rDPWwfLLGe0XIZ0E1ibFFjcUtZw=="},"inclusionProof":{"logIndex":"3199351599","rootHash":"xWGbayn3j6X2p0FV88LJRP73da2aNVUmQWGszpJE2wo=","treeSize":"3199351603","hashes":["YQwxpEIEzSTvnkviHSauZKeCAxph58u3To0yBcpVBcE=","geZoEKGlFJ4ENa+m4U4x1kJhJQAKEF6dSNb5wPqVDoM=","WvmZ1VcvxsbK/aVXJk1UtNuO3WTQn6Zv6V/NmveKTUY=","I6uUsXPH/yU9gxNayIxmRT8SyogQY4uaQiOLF6cAPv4=","v7HlsN/XaAEp8x3PlNUocxwJFivVsMZaIjPsqSNYP2U=","F/YmOyMBU6+REtpx345WJ1ZWOoaPh6GXTqSK7gWWUWI=","xqRGIj7vKKnGXGhCxIxXy1I2FKxqmR2s2WTICdDM2+8=","pkUNmvOtKnFBM2KMw42O+V67fz6YzLuD4OSqzBsCopI=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199351603\nxWGbayn3j6X2p0FV88LJRP73da2aNVUmQWGszpJE2wo=\n\n— rekor.sigstore.dev wNI9ajBGAiEA9UQzYZnyju2dDdhn+5rYni6KdxHtGf+svKDndo3FtCoCIQDk4G2ZVnYTavAdsmdkH0jn/yxNcu/WD+kU+H8gx/X3ew==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3M2M3OGEwOTU2OTg5NTEwMTdiMDk4MzM4MTVmMGEwNTMzOTc2YzE2MTlmZGU5ZGViYjFhYzEzYTQ3YWNhOTkxIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUM3eDMrM25sdWlDSDRrRzh6VERpY2NjbGJLdi9Vd3FkM3djYnErWXhjRW53SWhBSVlKTEtwN3hONW5hbWRRSURrNjhlbXhWcEJ1TlBXOWYvT2QwOEovalJuYyIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KTFowRjNTVUpCWjBsVllqQk1jR3A0Y2tzd1RXRmhLMXBQUkdkVlp6WkxiMUp1V2tOSmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJGTVZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVVeFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVU0YzJwbGNqUnhOMU5FZEROMWJVaHBOV2x4Vm1KV1ozWnNVSEZuUkZseFFqZHpURzhLTkN0NlVYZE9kemhoWjBKU00zZE9aMmx5VG05RVdtTlRkMFowVURVemFXNVdkMmhWWVc1RlNXRkhVbU52WkZsSVlYRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWUmNHVm5RVUZDUVUxQlVucENSa0ZwUlVGMFpXNUxkakpWYjFSMVQyUlFRMnB2ZERZd1NGVkJRbEU1VmtkNlZtY3dhZ3BVTDA1dlJrSXhiVTB3V1VOSlJDc3llakE0VldRMWVXVldZbW80YVhCT1RuVXdNelJsVmpsNVZHbExjbmQyVFZSWlJFZ3dlbEpZUWsxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKblFVMUhWVU5OUkVSNVFVVkNZMlU0WmxCeGQwZHpORlJTZFZsUlZrc3ZiR0lyWlVkRFdrOTFORXgyVDBacE9FSnphVXRKUldzS1VrMVpVVVV4UW5oS2VEZENSVzByZDFKUlNYaEJURWhTTXpGSmJtRlhZbXhOYlZOYWRVNXZWR0YwVFVjNVZFSlhZMUZCYWxnMFoyWTRSVVp4UldNemRncG1ZbkZDVG1JMVpqRkJUMnRzYkdGd1VXSllNV1JuUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgaijGhPg+orjO2HqFepxjI+9rHAGZFvdQQIxGtlIAFnwCFQDX44WaG5SY1RAX3VVDzzf3aQdzZRgPMjAyNjEwMTkxMzA1MTVaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTEzMDUxNVowLwYJKoZIhvcNAQkEMSIEIPOnevDgQakBFi1JLwObLE5zKHb842tollPIv6Pb/gltMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQD6HYpLK51sC96NU00okVJv8eKLCD9nAiG1FgFT4m784XtJh5ORXORCdDdTwx13lI8CMQCQJGY4Uf2ZPWKyqYEiZ+l03oBYThS6LXNPAWEmNuG+BefrJksySUIUHMxIHI3xnHI="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"c8eKCVaYlRAXsJgzgV8KBTOXbBYZ/eneuxrBOkesqZE="},"signature":"MEYCIQC7x3+3nluiCH4kG8zTDiccclbKv/Uwqd3wcbq+YxcEnwIhAIYJLKp7xN5namdQIDk68emxVpBuNPW9f/Od08J/jRnc"}} \ No newline at end of file diff --git a/build/torch210-cu130-aarch64-linux/symbols.json b/build/torch210-cu130-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch210-cu130-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch211-cu126-aarch64-linux/__init__.py b/build/torch211-cu126-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch211-cu126-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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-cu126-aarch64-linux/_ops.py b/build/torch211-cu126-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch211-cu126-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch211-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch211-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..10bf8eb7b6afadf1a91659397dbb79101a99b167 --- /dev/null +++ b/build/torch211-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b673ef243f23d746380c376309669b1d34f7fd7e615503cf070394e2962aea02 +size 286296 diff --git a/build/torch211-cu126-aarch64-linux/metadata.json b/build/torch211-cu126-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..5915ce3b1e6319c120416ff32b9e5be0fb632895 --- /dev/null +++ b/build/torch211-cu126-aarch64-linux/metadata.json @@ -0,0 +1,37 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "tnPvJD8j10Y4DDdjCWabHTT3/X5hVQPPBwOU4pYq6gI=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch211-cu126-aarch64-linux/metadata.json.sigstore b/build/torch211-cu126-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..07d90422af1a6c39bfbca09372feddcbef462096 --- /dev/null +++ b/build/torch211-cu126-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHKzCCBrGgAwIBAgIUEwXcKWEkaaKuxZNMxaUm0yXOi0wwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTE2WhcNMjYxMDE5MTMxNTE2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEZTZpwLai/krIrhQrDgQgJuf9HfvednhdVMWduVhdh13EW6NPYhf1eouJVmlttc54qZnKSnXAeKGh1V6AcRPQ4KOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVEQssAAABAMARjBEAiB/+DFzamA/TwmevX8+FQtqYcGo/XCBO7o5UXHhpjOipQIgU66h+EzQR9QtvzivaZmeDnM4bfCLGKZMKBnk/yPqJ+wwCgYIKoZIzj0EAwMDaAAwZQIwWhPfpdPAC3/6HJGkI1fuR/6u0x1mKUwgi3OWLQ7XZ3dlln1kDljSZ0CvCwaIW/TlAjEA7TX1K/CZV86w+qxCW1YetRnf6QIywTH80Zb8BpZ6lClChuB5oCG5LqzxiI1gDaG6"},"tlogEntries":[{"logIndex":"3321255916","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415116","inclusionPromise":{"signedEntryTimestamp":"MEUCIQC6odprucdqq+J8RWn7FaCV7oq+Ejw/gCCwHqm2UCN8SwIgIAcIwNquzYPF6igaWIUkKCYc/hSl+IFwqOgiP5e/bBM="},"inclusionProof":{"logIndex":"3199351654","rootHash":"ptJPGapd1rECzcoUrPxpmIKvfQ3EP3eibHw6sjj9jko=","treeSize":"3199351674","hashes":["whenpggRNyEMdbYCLMIXPIUkrboiS/1JcXIlKh6Z+Q0=","VxlfKxabBlpTngrlS4c9PTAMpvbmsk/o4hYkaRkpWlw=","UHdeZYaDEggatONM/kU9r0BfEz0RSTY1SrwYl9ShFyM=","29JqhnDa0k4MNcg+Vaxfh30ecmtVtgw2xkz26vewp6s=","/CB/z6PiogMZZDcjyElj6gik79FNhwxrPprAu8mN9EE=","wy5AHJ6RaEF2Dcyg4uwHfsgvXo3IikEvPpiSraqHGfo=","Zy4CynuBn7d4UhtJsu45MhTO6/W32PdogOi1xjoSMZY=","xqRGIj7vKKnGXGhCxIxXy1I2FKxqmR2s2WTICdDM2+8=","pkUNmvOtKnFBM2KMw42O+V67fz6YzLuD4OSqzBsCopI=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199351674\nptJPGapd1rECzcoUrPxpmIKvfQ3EP3eibHw6sjj9jko=\n\n— rekor.sigstore.dev wNI9ajBEAiBvtmhxvjOrGBfHAoiOqlqIotVHqL0V5begOIZXvvAGeAIgeFBnNuPSp+ldkXEAuBSeujBijy/xA4yNLJI4YcBI5ns=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIyYzMwYTJjZjIzOGRmNWQ2Y2ZjOWE2Yzg2MjA3ZGE4ZGUxMjRjNzRiYjZmYWQwMjc2NGM0M2NlMWQ0YTJjMzgwIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUURIU1QrRHhwUmk1RUtEMVJHNzVEc0lsNW4rU3R4TTlBYXQ5QVhQbnNzcHFnSWdJRzkrRElIaGJ2OHJQd1ViTFF1Tlp0bFJWcTd3UkZEYm9oam1aY1pjVk1jPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExla05EUW5KSFowRjNTVUpCWjBsVlJYZFlZMHRYUld0aFlVdDFlRnBPVFhoaFZXMHdlVmhQYVRCM2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJGTWxkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVVeVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZhVkZwd2QweGhhUzlyY2tseWFGRnlSR2RSWjBwMVpqbElablpsWkc1b1pGWk5WMlFLZFZab1pHZ3hNMFZYTms1UVdXaG1NV1Z2ZFVwV2JXeDBkR00xTkhGYWJrdFRibGhCWlV0SGFERldOa0ZqVWxCUk5FdFBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWUmMzTkJRVUZDUVUxQlVtcENSVUZwUWk4clJFWjZZVzFCTDFSM2JXVjJXRGdyUmxGMGNWbGpSMjh2V0VOQ1R6ZHZOUXBWV0Vob2NHcFBhWEJSU1dkVk5qWm9LMFY2VVZJNVVYUjJlbWwyWVZwdFpVUnVUVFJpWmtOTVIwdGFUVXRDYm1zdmVWQnhTaXQzZDBObldVbExiMXBKQ25wcU1FVkJkMDFFWVVGQmQxcFJTWGRYYUZCbWNHUlFRVU16THpaSVNrZHJTVEZtZFZJdk5uVXdlREZ0UzFWM1oya3pUMWRNVVRkWVdqTmtiR3h1TVdzS1JHeHFVMW93UTNaRGQyRkpWeTlVYkVGcVJVRTNWRmd4U3k5RFdsWTRObmNyY1hoRFZ6RlpaWFJTYm1ZMlVVbDVkMVJJT0RCYVlqaENjRm8yYkVOc1F3cG9kVUkxYjBOSE5VeHhlbmhwU1RGblJHRkhOZ290TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgUh2FSVZoBrIz1juXWq/2Gc9O2VqDF7KsJZD+xIb5qEoCFQDozEX0CI0Ekt8DqmNYZwkoJr6EGBgPMjAyNjEwMTkxMzA1MTZaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTEzMDUxNlowLwYJKoZIhvcNAQkEMSIEIPp/o44Uo4mEnoQuNnndq6swDHR7cmty9yn1cGSyDqizMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQCCkqTgqy88KRGkB87Q4ZSIv1ykqLia1YtaLfqsDSlLWJzXdt97iMmfy1qhYfHEItsCMQClKlCxeehD5lAtB8IO3S3oWZ5H7lEhkvbL4tMSskKJF0+3Aj8A/Y5zdBJEEqP73iY="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"LDCizyON9dbPyabIYgfajeEkx0u2+tAnZMQ84dSiw4A="},"signature":"MEUCIQDHST+DxpRi5EKD1RG75DsIl5n+StxM9Aat9AXPnsspqgIgIG9+DIHhbv8rPwUbLQuNZtlRVq7wRFDbohjmZcZcVMc="}} \ No newline at end of file diff --git a/build/torch211-cu126-aarch64-linux/symbols.json b/build/torch211-cu126-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch211-cu126-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch211-cu128-aarch64-linux/__init__.py b/build/torch211-cu128-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch211-cu128-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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-aarch64-linux/_ops.py b/build/torch211-cu128-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch211-cu128-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch211-cu128-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch211-cu128-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..db04f7e24fd341b657b2df2d7c7a0d6ce901a561 --- /dev/null +++ b/build/torch211-cu128-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dee43e47efb0108f66eab8017cab7fa9c76b233a22e64e92bef41f4a785f643e +size 417376 diff --git a/build/torch211-cu128-aarch64-linux/metadata.json b/build/torch211-cu128-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..96f31247aca6ee19b0c31e3698f8b5f807de3e18 --- /dev/null +++ b/build/torch211-cu128-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "3uQ+R++wEI9m6rgBfKt/qcdrIzoi5k6SvvQfSnhfZD4=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch211-cu128-aarch64-linux/metadata.json.sigstore b/build/torch211-cu128-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..b77786f8de4518f1dd7ac74d73bcf9bf15c2a1cc --- /dev/null +++ b/build/torch211-cu128-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrKgAwIBAgIUYuzSNv2HITt2l5TLDyt/ZIh+CwMwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTE3WhcNMjYxMDE5MTMxNTE3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEUOUtvCTDTPiGioFa4xqs8K4kwFr2AltlGR1CS0q1G/2vA9Zc8l/VLSR6yw5419KQTOaQhcgRkYyCgNJJjDuPv6OCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVEQwlQAABAMARzBFAiEA+FmHocOssmxnmXuhoL7Y6s+jlB6vgl89A1CLdV9qlfACIHVwzXitpB63q1gUTxISR8xHFyZOvt5CtQAcSipFOyW0MAoGCCqGSM49BAMDA2gAMGUCMDb/GVaanUHoYae1miaVDmxkTz683X/yHEkw4yD3UBbC7gX1u8hiF53QwG85nC5nOAIxAPjvjuMmOxp9BdP6b0rjIKviOINSFb6ImAnfn5TThuAsFcX+txLfhH5laAfC6eV1tg=="},"tlogEntries":[{"logIndex":"3321256015","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415117","inclusionPromise":{"signedEntryTimestamp":"MEUCICe5mxcuapXYPv/o4zF4KPnPUD2eiHFJt6UPE4YuWZBRAiEArIvyCPWXmPlpVBQbgFxCNrPTNJOysdFHC9EQmO3gxm0="},"inclusionProof":{"logIndex":"3199351753","rootHash":"WfsTKjrXn0cAdcVo7L21IUYplE5qoOmAEDMVfP3X4pQ=","treeSize":"3199351779","hashes":["bj32ptHqUpMw8pQm9CVVHjISMJSV8FURatKtDilzWmA=","f3KmMTOcRe+3psvjm8fptg+RR6+HWu51MlLL6hQ3/9I=","TjUbCekDQwF9yyZgNtDoR9lSTE7QvH1zvNlHPqeVauo=","REdYxVA7qeDZPEJEjyJmmSTRRbl10xn00lbbqeyokgw=","km2gYFbJNPVnwpafqpSQNTZvbN+kZ3s5kwLplMDQk1s=","Et5KoZ6Kay6wxd5J7+o3p8RrUURy4N9GA7vgOJPV1gk=","xnu3GBBam9eIEFj1uWV4MzenQZ9dv+Q86dm2IFcBc5s=","UDKQzdNzI/5wP5RT9aCwOlCtRShDeXkPRX2KJPlQQlo=","xqRGIj7vKKnGXGhCxIxXy1I2FKxqmR2s2WTICdDM2+8=","pkUNmvOtKnFBM2KMw42O+V67fz6YzLuD4OSqzBsCopI=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199351779\nWfsTKjrXn0cAdcVo7L21IUYplE5qoOmAEDMVfP3X4pQ=\n\n— rekor.sigstore.dev wNI9ajBFAiAt+XJ1CXhjS4ECXe61eQzHT/3vSnk3dNKWBmCtuwyS3QIhAPiIhqbTEEnTmXf89W1Ey1H2KqAeOYD33EGDIbRkqahA\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI0MTc5YThhNDczOGU4MWU5NjllOTJkMjZkY2U5NjQ4YTI5NDdlOTJkNDEwMjNlNzJmZjM1YzJhZWVlZTc4NjBjIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJQnNsV3UyR1FqZmRHenM4S2xUclFJQ2RNVS9UOXJDVC9KbmlmQzh0dzNOT0FpQjVualVkaTRUTWNZbW1TbkI5NGVWK2ladExlUFIwVWpIU2dENHE1MlBKVUE9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KTFowRjNTVUpCWjBsVldYVjZVMDUyTWtoSlZIUXliRFZVVEVSNWRDOWFTV2dyUTNkTmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJGTTFkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVVelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZWVDFWMGRrTlVSRlJRYVVkcGIwWmhOSGh4Y3poTE5HdDNSbkl5UVd4MGJFZFNNVU1LVXpCeE1VY3ZNblpCT1Zwak9Hd3ZWa3hUVWpaNWR6VTBNVGxMVVZSUFlWRm9ZMmRTYTFsNVEyZE9Ta3BxUkhWUWRqWlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWUmQyeFJRVUZDUVUxQlVucENSa0ZwUlVFclJtMUliMk5QYzNOdGVHNXRXSFZvYjB3M1dUWnpLMnBzUWpaMloydzRPUXBCTVVOTVpGWTVjV3htUVVOSlNGWjNlbGhwZEhCQ05qTnhNV2RWVkhoSlUxSTRlRWhHZVZwUGRuUTFRM1JSUVdOVGFYQkdUM2xYTUUxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKblFVMUhWVU5OUkdJdlIxWmhZVzVWU0c5WllXVXhiV2xoVmtSdGVHdFVlalk0TTFndmVVaEZhM2MwZVVRelZVSmlRemRuV0RFS2RUaG9hVVkxTTFGM1J6ZzFia00xYms5QlNYaEJVR3AyYW5WTmJVOTRjRGxDWkZBMllqQnlha2xMZG1sUFNVNVRSbUkyU1cxQmJtWnVOVlJVYUhWQmN3cEdZMWdyZEhoTVptaElOV3hoUVdaRE5tVldNWFJuUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgPaoP9+6j9HfULwCp+WXxqGPcZqmLYiKtOH0yQaDRYkECFHRF/gJRec9XXQ5vDke7KGUJueQLGA8yMDI2MTAxOTEzMDUxN1owAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTMwNTE3WjAvBgkqhkiG9w0BCQQxIgQgvNSvrZ7QMFGCN46cYZ6l+PuxwELXOMMALfIFDk/xLpEwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAKFd37IIKn3fQuWDzyNlIbz9kJTui8E/fa+iQyrN/q1zWoH5SxAdAUc6hlhZhhXgvwIxALZ2sMuow9s0hYNXPpAMIOY5KtZ1qBTvtKlAWFpwoE1CHhUPaZpj5ZqpehezdMyDoQ=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"QXmopHOOgelp6S0m3OlkiilH6S1BAj5y/zXCru7nhgw="},"signature":"MEQCIBslWu2GQjfdGzs8KlTrQICdMU/T9rCT/JnifC8tw3NOAiB5njUdi4TMcYmmSnB94eV+iZtLePR0UjHSgD4q52PJUA=="}} \ No newline at end of file diff --git a/build/torch211-cu128-aarch64-linux/symbols.json b/build/torch211-cu128-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch211-cu128-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch211-cu130-aarch64-linux/__init__.py b/build/torch211-cu130-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch211-cu130-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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-cu130-aarch64-linux/_ops.py b/build/torch211-cu130-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch211-cu130-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch211-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch211-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..3bf554ec3e7191e9b9acb4112906ac892b041a5b --- /dev/null +++ b/build/torch211-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0191b655026e7f4b6c43fb8253fa37d45b431db3e122e00bc9926a2190b43920 +size 421472 diff --git a/build/torch211-cu130-aarch64-linux/metadata.json b/build/torch211-cu130-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..19e3184d945e0ea86b9ba09a282af23e10e336f1 --- /dev/null +++ b/build/torch211-cu130-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "AZG2VQJuf0tsQ/uCU/o31FtDHbPhIuALyZJqIZC0OSA=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch211-cu130-aarch64-linux/metadata.json.sigstore b/build/torch211-cu130-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..264ed3529f4ce4f8d1168c73f6b9af93272c422b --- /dev/null +++ b/build/torch211-cu130-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrKgAwIBAgIUQesPOCKrj6K5/hogXsnytdrtqt8wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTE4WhcNMjYxMDE5MTMxNTE4WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEn6bZ2wujSgZM4tVqE+V5WCForPGi33dQxm2Z5B1BzEvjQchJ51kk1TomHOLBwCPn6//RWZzJtiicsssXsY3FLaOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVEQ0pwAABAMARzBFAiEAp8zKk+5Uxm4/3w458/n4mVXrnPeRFGTd+g3Ari8ikhkCIEm56I50ilb25QdvRQRUt00aP6f3Lz420qluSkOR3ZQ1MAoGCCqGSM49BAMDA2gAMGUCMBpXONFElaAL4oNo41QcFQUSzvGa9Fk65V03zEFSIGA0RUgSwTdvGzrTALA1s52nnAIxAMJ+niq7f8iRtgvDCKjnXADSHbc/nQC0qJj2acyab9fkoY1qG5gfPufpLaqp19/LuQ=="},"tlogEntries":[{"logIndex":"3321256127","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415118","inclusionPromise":{"signedEntryTimestamp":"MEQCIHGKRanhmhpIsY9vOKafGwzlbZRSeoSM/A2hEs3BknANAiAsfRKlOxhshrREYpSOA4yhNc27WIgrEMFtVCcvcUzDZg=="},"inclusionProof":{"logIndex":"3199351865","rootHash":"0OUvbmxd/M04TmKBs15Eftp4lZM5kofoauIq6PerwAc=","treeSize":"3199351868","hashes":["gqm7FjAuGat6WuPNeew1HTunTU8gvRw2xrNLSE9pdss=","qVYxXdJ2jdfakeWb+WBWhVPI0/mkt/2BL+4Yi8UlPIw=","Uxtevdz7BsmkoE+qRZzR2jQ5vt8suQCQD3HQ3umRRHQ=","SVu3R3XZUINk0SVDi5k3SMuXc7AqoaXm3r0a8lpfO8Y=","1C82OTkToDumAKk+kgLxoWAAfylqvbbutGXC4kJ5XFs=","pvSOMtCoDY7VfblJPv4FHsSUD1IfWr1b9YUvy3C4W64=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199351868\n0OUvbmxd/M04TmKBs15Eftp4lZM5kofoauIq6PerwAc=\n\n— rekor.sigstore.dev wNI9ajBGAiEAzhYuQgYDv/IBrgYJMwbK3IrppjchBweJx8dugXOfUe0CIQCFxVROEnt+WEd+olCjIRWEwQ0YNlBNDCS9VhJx3PUeuQ==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJmZThmODdmZGUwMWYxNjIyMjhiOTIzNTlkMzg0Njc0NzAzZWQwNDlmM2RiN2QwZGFjZjk4MGQxYzI2YTk3NDg4In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUQzZC9jWWlqZ2xyc25DMm1pQStIdGRtYm5qOWZhbVFyc1pPSmhEQ0FuRG1BSWdYei9TWmczeXU4bUZZMUprSDY3SW1WVjBkUEJMbjJmWUpvUTJlbjlXeTBJPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KTFowRjNTVUpCWjBsVlVXVnpVRTlEUzNKcU5rczFMMmh2WjFoemJubDBaSEowY1hRNGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJGTkZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVVMFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ1Tm1KYU1uZDFhbE5uV2swMGRGWnhSU3RXTlZkRFJtOXlVRWRwTXpOa1VYaHRNbG9LTlVJeFFucEZkbXBSWTJoS05URnJhekZVYjIxSVQweENkME5RYmpZdkwxSlhXbnBLZEdscFkzTnpjMWh6V1ROR1RHRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWUk1IQjNRVUZDUVUxQlVucENSa0ZwUlVGd09IcExheXMxVlhodE5DOHpkelExT0M5dU5HMVdXSEp1VUdWU1JrZFVaQW9yWnpOQmNtazRhV3RvYTBOSlJXMDFOa2sxTUdsc1lqSTFVV1IyVWxGU1ZYUXdNR0ZRTm1ZelRIbzBNakJ4YkhWVGEwOVNNMXBSTVUxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKblFVMUhWVU5OUW5CWVQwNUdSV3hoUVV3MGIwNXZOREZSWTBaUlZWTjZka2RoT1Vack5qVldNRE42UlVaVFNVZEJNRkpWWjFNS2QxUmtka2Q2Y2xSQlRFRXhjelV5Ym01QlNYaEJUVW9yYm1seE4yWTRhVkowWjNaRVEwdHFibGhCUkZOSVltTXZibEZETUhGS2FqSmhZM2xoWWpsbWF3cHZXVEZ4UnpWblpsQjFabkJNWVhGd01Ua3ZUSFZSUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgzvdyIdODeM0YnUoP+hgtXgYPcqcz1BC7duy9CRksdJMCFQCsyWbgRqhzuy/OZmvZebcBs3c5uhgPMjAyNjEwMTkxMzA1MThaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTEzMDUxOFowLwYJKoZIhvcNAQkEMSIEIHlOFZEefHUfGLzqqu6QvKECiAKwJHEabu280z/hBAqgMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQCEzPZ+N2gs5RngX4jI+FqspZIbKMXx1XSBlcbxrmcKsGvszg/L4kbwCpUt2AhE3JECMAy1LqP18jd3lZIOmdmRgNN4ev3G2r9mYdSx9E+fX22dYp/Nj1s+DwKjRbWljJom7w=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"/o+H/eAfFiIouSNZ04RnRwPtBJ89t9Daz5gNHCapdIg="},"signature":"MEUCIQD3d/cYijglrsnC2miA+Htdmbnj9famQrsZOJhDCAnDmAIgXz/SZg3yu8mFY1JkH67ImVV0dPBLn2fYJoQ2en9Wy0I="}} \ No newline at end of file diff --git a/build/torch211-cu130-aarch64-linux/symbols.json b/build/torch211-cu130-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch211-cu130-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch212-cu126-aarch64-linux/__init__.py b/build/torch212-cu126-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch212-cu126-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch212-cu126-aarch64-linux/_ops.py b/build/torch212-cu126-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch212-cu126-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch212-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch212-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..bb033a07ff82d219680080d675172cb2a3b25297 --- /dev/null +++ b/build/torch212-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8c5dceb4054e89677081abf0ded0a5db2cbfc7b943d3d11c67ecb7f0af0a4f1b +size 286088 diff --git a/build/torch212-cu126-aarch64-linux/metadata.json b/build/torch212-cu126-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..2560b3aa1d7f6ff761b081ddc095ae69b039551e --- /dev/null +++ b/build/torch212-cu126-aarch64-linux/metadata.json @@ -0,0 +1,37 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "jF3OtAVOiWdwgavw3tCl2yy/x7lD09EcZ+y38K8KTxs=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch212-cu126-aarch64-linux/metadata.json.sigstore b/build/torch212-cu126-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..3b9a549ef22d68a2f168e6a3861c80a9e2598aa6 --- /dev/null +++ b/build/torch212-cu126-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrKgAwIBAgIUKIPksPNUWSxGZ3XYp0TRkd/qtVkwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTE5WhcNMjYxMDE5MTMxNTE5WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEDkfyk3ghF14pDH+VEQMQei9DFLNhm9jOZ5XaCK6Gsy9A+cJaga/OkP4SBu0CECmrAzEJBfUXuHLYjMCZAT/7o6OCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVEQ37AAABAMARzBFAiEAz1pdtR4P0scI7byDos8eK/0LhS+NoxX4rDJZSh9tbw4CIElmrUA/8gbAUXSLZuFp3KXeQONj2gX4pamhYrCjkyJRMAoGCCqGSM49BAMDA2gAMGUCMFlz5KymFD494B9DVLjWJhPPtMfRSweGJb9Kg232nqcRQ7mq3jwUk60ErcG4JiS9rwIxAP7ZG3uoKDqeTtZ0ItGoOXptcceIKeP2cmZCzRZCapRnWlwa7eWN9NE1aD68lmqc7g=="},"tlogEntries":[{"logIndex":"3321256189","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415119","inclusionPromise":{"signedEntryTimestamp":"MEUCIFhLr4FHTIcLRViVybjrXhlgMMGsDOA5iuVMlhj7GHCfAiEAly/nZEqX9w4fJ1TXUZAGilRy7alWaoSdNN6tZUX1HJU="},"inclusionProof":{"logIndex":"3199351927","rootHash":"5u4xS7ehWXpHl0WB+xCs1MDgCj0r6jOi4pDvCXdzbRs=","treeSize":"3199351928","hashes":["jzeqVe4Mk+f+a+25AxqZHz3XXOXdrrnVGIh6YDGwL0Q=","EEsdDNVhz6leAMZDPMS81jeQtjdjhyXPU1xmQS5mbPA=","svuEruPG8rNgkFscAx9Y72tz5SrhSUgwlznV2eSkY9s=","7tOk1CAwe34+sVeO0uqFCP1LPNIWGrbRsGBYCe8KyBY=","4J18rxcO2EhjBfiExf1coS/BQOXPnnP4Q4gxS4nTJh0=","5V+Y6Y1T1j0sNR+jxfPl1dXF7MKqR63vqHvv7e9W4+U=","pvSOMtCoDY7VfblJPv4FHsSUD1IfWr1b9YUvy3C4W64=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199351928\n5u4xS7ehWXpHl0WB+xCs1MDgCj0r6jOi4pDvCXdzbRs=\n\n— rekor.sigstore.dev wNI9ajBFAiEA1s+tzznnnsvp9Fbq5+DeOSPHcS2UDHkYW21R9Fl5YXUCICN1pGfKoo14xVGUO4uGHUbOE/jmFfJGvtHt0tdPhSlM\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIzYjkyZjhhNDQ5MjAzNDQ0OWU5MTUzNGJkNGY0Mjg1ODUxMjBhMmRiNTllYmRkZjI3YjI5NDMwYzkwODFmZmQ0In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNMNGxyNzlEWUpuOCtRRVRPa21BSk1FMmhVeDVSU25FRTVWY1BxMUJhQUpBSWhBSWhkMkg2K040N0thMzcxRnNackErUTR4NVozOWRud2xSNFA1RCszVHFUcyIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KTFowRjNTVUpCWjBsVlMwbFFhM05RVGxWWFUzaEhXak5ZV1hBd1ZGSnJaQzl4ZEZacmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJGTlZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVVMVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZFYTJaNWF6Tm5hRVl4TkhCRVNDdFdSVkZOVVdWcE9VUkdURTVvYlRscVQxbzFXR0VLUTBzMlIzTjVPVUVyWTBwaFoyRXZUMnRRTkZOQ2RUQkRSVU50Y2tGNlJVcENabFZZZFVoTVdXcE5RMXBCVkM4M2J6WlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWUk16ZEJRVUZDUVUxQlVucENSa0ZwUlVGNk1YQmtkRkkwVURCelkwazNZbmxFYjNNNFpVc3ZNRXhvVXl0T2IzaFlOQXB5UkVwYVUyZzVkR0ozTkVOSlJXeHRjbFZCTHpobllrRlZXRk5NV25WR2NETkxXR1ZSVDA1cU1tZFlOSEJoYldoWmNrTnFhM2xLVWsxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKblFVMUhWVU5OUm14Nk5VdDViVVpFTkRrMFFqbEVWa3hxVjBwb1VGQjBUV1pTVTNkbFIwcGlPVXRuTWpNeWJuRmpVbEUzYlhFS00ycDNWV3MyTUVWeVkwYzBTbWxUT1hKM1NYaEJVRGRhUnpOMWIwdEVjV1ZVZEZvd1NYUkhiMDlZY0hSalkyVkpTMlZRTW1OdFdrTjZVbHBEWVhCU2JncFhiSGRoTjJWWFRqbE9SVEZoUkRZNGJHMXhZemRuUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg6r4mSSbLZ33fkYsaWbAitOPUp7tlkOsdORRvUggnSW0CFA7K//cen5dSaEc5nLWF52WZMmNxGA8yMDI2MTAxOTEzMDUxOVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTMwNTE5WjAvBgkqhkiG9w0BCQQxIgQgk8RvT+anXR0P7uu8hGuB4D8cG1NT/M9GdP5PKMO35pUwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAJBfa73BTX2sJWbwnqditSsoMzFGDTvHGUEY3H/18j0PfAE13675a7W9hhhZas7AGQIxAOQIcKUnjv0kV5xfwcU7EZZgCqM17DPWjRI5t4gyGBIj094c/robZj0RrZdLf4JV2g=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"O5L4pEkgNESekVNL1PQoWFEgottZ693yeylDDJCB/9Q="},"signature":"MEYCIQCL4lr79DYJn8+QETOkmAJME2hUx5RSnEE5VcPq1BaAJAIhAIhd2H6+N47Ka371FsZrA+Q4x5Z39dnwlR4P5D+3TqTs"}} \ No newline at end of file diff --git a/build/torch212-cu126-aarch64-linux/symbols.json b/build/torch212-cu126-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch212-cu126-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch212-cu130-aarch64-linux/__init__.py b/build/torch212-cu130-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch212-cu130-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch212-cu130-aarch64-linux/_ops.py b/build/torch212-cu130-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch212-cu130-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch212-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch212-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..639f038d7a8a8e695547dede637b8dc1fa4b527a --- /dev/null +++ b/build/torch212-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7c342bbc052eee8c403763d44c3117eece563fe92306febc080c29e442277d17 +size 421408 diff --git a/build/torch212-cu130-aarch64-linux/metadata.json b/build/torch212-cu130-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..64909f4e0fd776038f6f85f9562416f6a1b42196 --- /dev/null +++ b/build/torch212-cu130-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "fDQrvAUu7oxAN2PUTDEX7s5WP+kjBv68CAwp5EInfRc=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch212-cu130-aarch64-linux/metadata.json.sigstore b/build/torch212-cu130-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..c0b559995f71974b155230c5e6d1b46d88b8cdde --- /dev/null +++ b/build/torch212-cu130-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrKgAwIBAgIUdQqhQHv7PcVJwP9MfX8QSs/cIgswCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTIwWhcNMjYxMDE5MTMxNTIwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEndnXPxkD8Oj/kfFHHHFIFPq7FXxFH6uk940egpSJuS0oTgH64qKgoaJ6vSSJQkd9wX2ZtDyYXzjec1YneyeNg6OCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVEQ8RAAABAMARzBFAiEAn/uZ6YP5wRFCub5WBswscpccoggoz5XTslU56Pd+FRUCIGdYHpSv2iXUFyvbYaaS8b3j0N8DDWeI21Vs2A9Ub0CWMAoGCCqGSM49BAMDA2gAMGUCMDgwfor+nEL7SKHFZnAKOWxKVmHyX1PNd2oM9YXxcskboyP8k20sLW/6mDFaYXvgogIxAJhMp0YjaTyzUY6W48AH3TUhC2Ipa0dFXoYZDchTjrzU6Va04GaAzLmd0bpe7nO+AQ=="},"tlogEntries":[{"logIndex":"3321256304","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415120","inclusionPromise":{"signedEntryTimestamp":"MEUCIQChxzihL88jHh9baI3rJhbbYPxJNW6PoPCtUYzbwF/obAIgflLNuZt7vh/vwCzjDy+c0pzUlxOiFNbniigFD6uO/eY="},"inclusionProof":{"logIndex":"3199352042","rootHash":"cpRiEJzYyzg2Lvl+tW4/f+RBeCtMuQBdIm9focmZik8=","treeSize":"3199352048","hashes":["IfSBlbsFDXzre5Qld0zyK7cUguFqGRZ6lZMSZvA7A6A=","c5PxnPn09FkoQjQn+xsQbPYxbwhEYM8d67wTkCY1kRo=","NWRdP1vLN48ASGLYRVAhDjLLFUa7G/UKV5iAssFeNZQ=","BQ/oUgeOEh9zjB+Rus/uvUWm4HEURFVrmPoaQ0xclq8=","6PTjbrgF4gmJMXQD4N5eSW/oLOOeuwXfLZBnah+9hJ8=","/r1MjldhvTi5OhOJCp+Dx1xd4N0f5cyGF1KLy8wuWrY=","ig9pr77VIw1H//BUEywcKR6qnlwMqWgbOLe5HSB+RtQ=","pvSOMtCoDY7VfblJPv4FHsSUD1IfWr1b9YUvy3C4W64=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199352048\ncpRiEJzYyzg2Lvl+tW4/f+RBeCtMuQBdIm9focmZik8=\n\n— rekor.sigstore.dev wNI9ajBEAiBtOuFl9CsT9E4lIEmGEXLs29tldrNMX7vpUXYN1wOWDgIgChsceXiFp1MHPzS34dN1e34U0Ov6HvG5JE8pj1rxbfM=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI1YjY1Zjk2YThjNjg1NzlhZmQzZjc1MWE2ZGZiODM4NTMyYjc3OTA4ZThjZWU4OTgwMzYyYjk5NmI3MTVjOGQ2In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUURFZkpvKzVUVEZHU0pRNlZZWXZ0c3o2dWZGdFJPYVk5cGM4anY1T1FXNFJ3SWdNeWIwNDV0OVNrZGtYZ3QxZUVFOVpBZ1JudUY5ZTBGR2VQOHdiWGlwa3RvPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KTFowRjNTVUpCWjBsVlpGRnhhRkZJZGpkUVkxWktkMUE1VFdaWU9GRlRjeTlqU1dkemQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJKZDFkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVsM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ1Wkc1WVVIaHJSRGhQYWk5clprWklTRWhHU1VaUWNUZEdXSGhHU0RaMWF6azBNR1VLWjNCVFNuVlRNRzlVWjBnMk5IRkxaMjloU2paMlUxTktVV3RrT1hkWU1scDBSSGxaV0hwcVpXTXhXVzVsZVdWT1p6WlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWUk9GSkJRVUZDUVUxQlVucENSa0ZwUlVGdUwzVmFObGxRTlhkU1JrTjFZalZYUW5OM2MyTndZMk52WjJkdmVqVllWQXB6YkZVMU5sQmtLMFpTVlVOSlIyUlpTSEJUZGpKcFdGVkdlWFppV1dGaFV6aGlNMm93VGpoRVJGZGxTVEl4Vm5NeVFUbFZZakJEVjAxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKblFVMUhWVU5OUkdkM1ptOXlLMjVGVERkVFMwaEdXbTVCUzA5WGVFdFdiVWg1V0RGUVRtUXliMDA1V1ZoNFkzTnJZbTk1VURnS2F6SXdjMHhYTHpadFJFWmhXVmgyWjI5blNYaEJTbWhOY0RCWmFtRlVlWHBWV1RaWE5EaEJTRE5VVldoRE1rbHdZVEJrUmxodldWcEVZMmhVYW5KNlZRbzJWbUV3TkVkaFFYcE1iV1F3WW5CbE4yNVBLMEZSUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgAvMbpwWt3G3+jyJMmXWl3Gt1VvMa9z26rzjwZ/CrOXYCFQDwoPEPmiz8IQtagQ/N+YbfQ/mSbhgPMjAyNjEwMTkxMzA1MjBaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTEzMDUyMFowLwYJKoZIhvcNAQkEMSIEIC6g4WWHYGQAK1A+/Ai0anditASHIVyryEaJiZsaLBfyMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMF/+Ar94KnFfRyJWLGQC7sd+iv+BurTjc3SbGT6aJSiRddH7uIbYbwGdTJDqc+RINwIwaewaIg7JxSAV0bOdVE4ylUY0w1isL83NWhmcacRZRRlP4d90TaNY7kHeaCXtpbIb"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"W2X5aoxoV5r9P3UabfuDhTK3eQjozuiYA2K5lrcVyNY="},"signature":"MEUCIQDEfJo+5TTFGSJQ6VYYvtsz6ufFtROaY9pc8jv5OQW4RwIgMyb045t9SkdkXgt1eEE9ZAgRnuF9e0FGeP8wbXipkto="}} \ No newline at end of file diff --git a/build/torch212-cu130-aarch64-linux/symbols.json b/build/torch212-cu130-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch212-cu130-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch212-cu132-aarch64-linux/__init__.py b/build/torch212-cu132-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch212-cu132-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch212-cu132-aarch64-linux/_ops.py b/build/torch212-cu132-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch212-cu132-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch212-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch212-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..11f3e3001e8e959ac1f6b9021d78c2552bac38b7 --- /dev/null +++ b/build/torch212-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:685be0470bab466b205301c70e3f083b966671c41bbd8d1688550006a0b90294 +size 489256 diff --git a/build/torch212-cu132-aarch64-linux/metadata.json b/build/torch212-cu132-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..b498d54e57b0b64e96d166f23f81ed73d1cde3be --- /dev/null +++ b/build/torch212-cu132-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "aFvgRwurRmsgUwHHDj8IO5ZmccQbvY0WiFUABqC5ApQ=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch212-cu132-aarch64-linux/metadata.json.sigstore b/build/torch212-cu132-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..8309826742907ac81fff6764613171379a439986 --- /dev/null +++ b/build/torch212-cu132-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHKzCCBrKgAwIBAgIUFtpPL7tQGF8T+k1XnVPh7M9S9wIwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTIxWhcNMjYxMDE5MTMxNTIxWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEuG04DMCFZYcYVkyFTwKVAo58KGT72Dtw66YtNO1iwWJTIjV86osQ3vUyuRdL4gRe8TMJ6NWGCgjfNPhlC7GBHKOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVEQ/3wAABAMARzBFAiEAmFEI3sbt8Yrd70M/eafO+4nIhgxsQtcfFCfsNVoEOyYCIF6QINzMlC56xHpo2YpiHIYMRJPa1ShtfRiK2/vxHB85MAoGCCqGSM49BAMDA2cAMGQCMBsa8IVo1495JXKbnpzGcT2KWATJn/c0jCmQfbCWrK+ITci54z8Ui3uCYxFGn4MJ+gIwOy3hoUBtIgucfiagTsG30TYbn39oiSdjAkes74ZqMhx4xwpFSfwY4314DGrHa1uK"},"tlogEntries":[{"logIndex":"3321256412","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415121","inclusionPromise":{"signedEntryTimestamp":"MEUCIANn7ZtvTHXSZFyIcyu0CRyCRXKc03HG5D2g1naUxU7cAiEAuDWQPytDbDKQ7kZEPb+oWtXPQoYUbMImc3Gek3/SWdQ="},"inclusionProof":{"logIndex":"3199352150","rootHash":"O5ScmTT57omaY7PjaDFz7+94nLGCf9FPPbbzjdDr2OU=","treeSize":"3199352154","hashes":["EAT+TK2x8TrmQJcbNBFmVCh5vIqyglcIXIk7rMCps0k=","zJLbV1oX8Wn1w7JaxNMLwsnZHxpMhgZtr04KZ01W4Fc=","jLEsIITHwHiGbX/J1rOxzQDg43+lo1dbr5CoYCf913s=","QxY+8BXJUbUORzw9TvQhwESHcI70jc7yHyUotiZez8s=","NgCW8eTVKpdotQRTwNksOGvAWDzMPXkJSquzaGW6FPk=","DNVD8/ILD01vey0jQKVj0dFyM37IO0xtg183GkIknEY=","zRydiKTQ2lwubTShI6BKXR2DLH/300r/vwdstdu9bFs=","pvSOMtCoDY7VfblJPv4FHsSUD1IfWr1b9YUvy3C4W64=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199352154\nO5ScmTT57omaY7PjaDFz7+94nLGCf9FPPbbzjdDr2OU=\n\n— rekor.sigstore.dev wNI9ajBFAiBym+qhpL1UytLyDF5xCKORZqp2h0NTjhPol+IFAcwHXgIhAJQeoe7aTcqfRpkS9DK/pSFcwMSbN2lKLFeCWXINU3uV\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJmZjEyZTA2NTVmNmY2NzkxNjgyZTI4MDMzMTQ5NWYwZmE1MTZjNDc1OTFmODFjMmM0Mjc0ZjNlZjBmMjFkMWY5In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJSEVES0RkYnFOSGJzeERZaFUrZmJOa3hDVTE1amVhVXZhZWJtKzhuSWNhZEFpRUFrbytxcjJCdFF4SktndXdHNUJEeFg2cW53Q0orYXYzcTVxMFRvTGVQRUs0PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExla05EUW5KTFowRjNTVUpCWjBsVlJuUndVRXczZEZGSFJqaFVLMnN4V0c1V1VHZzNUVGxUT1hkSmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJKZUZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVsNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVYxUnpBMFJFMURSbHBaWTFsV2EzbEdWSGRMVmtGdk5UaExSMVEzTWtSMGR6WTJXWFFLVGs4eGFYZFhTbFJKYWxZNE5tOXpVVE4yVlhsMVVtUk1OR2RTWlRoVVRVbzJUbGRIUTJkcVprNVFhR3hETjBkQ1NFdFBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWUkx6TjNRVUZDUVUxQlVucENSa0ZwUlVGdFJrVkpNM05pZERoWmNtUTNNRTB2WldGbVR5czBia2xvWjNoelVYUmpaZ3BHUTJaelRsWnZSVTk1V1VOSlJqWlJTVTU2VFd4RE5UWjRTSEJ2TWxsd2FVaEpXVTFTU2xCaE1WTm9kR1pTYVVzeUwzWjRTRUk0TlUxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKalFVMUhVVU5OUW5OaE9FbFdiekUwT1RWS1dFdGlibkI2UjJOVU1rdFhRVlJLYmk5ak1HcERiVkZtWWtOWGNrc3JTVlJqYVRVS05IbzRWV2t6ZFVOWmVFWkhialJOU2l0blNYZFBlVE5vYjFWQ2RFbG5kV05tYVdGblZITkhNekJVV1dKdU16bHZhVk5rYWtGclpYTTNORnB4VFdoNE5BcDRkM0JHVTJaM1dUUXpNVFJFUjNKSVlURjFTd290TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgJoabU4djISjj7nouh6WiGnraOEJm6g5oTINyM/aqi0YCFB4XTCzBJftoSsk1pNs26PVGzQk4GA8yMDI2MTAxOTEzMDUyMVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTMwNTIxWjAvBgkqhkiG9w0BCQQxIgQg3qE67WS4ktZbrzLqs61aXlONiPVZ/T7WiVZwYSMz0l4wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxALIum9Y7vxavTWu+UZp5i8rt0uUe6fHWUKLk1ILje03Y1WY8SuGyxVCCsshjzDDv0AIwZCPE+sAfcDrp2PklQNx7xoo24V5YngLgOo5CgXmdieDIPp7Ha/b8uDJyYUJUOmOV"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"/xLgZV9vZ5FoLigDMUlfD6UWxHWR+BwsQnTz7w8h0fk="},"signature":"MEUCIHEDKDdbqNHbsxDYhU+fbNkxCU15jeaUvaebm+8nIcadAiEAko+qr2BtQxJKguwG5BDxX6qnwCJ+av3q5q0ToLePEK4="}} \ No newline at end of file diff --git a/build/torch212-cu132-aarch64-linux/symbols.json b/build/torch212-cu132-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch212-cu132-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch213-cu126-aarch64-linux/__init__.py b/build/torch213-cu126-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch213-cu126-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch213-cu126-aarch64-linux/_ops.py b/build/torch213-cu126-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch213-cu126-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch213-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch213-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..4bce9377097ba17396871a5e9828312c14771e71 --- /dev/null +++ b/build/torch213-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5bd28d1c4dc261550b789138b2a55cd70f81454701b658f68b3427d2e835cb5c +size 286088 diff --git a/build/torch213-cu126-aarch64-linux/metadata.json b/build/torch213-cu126-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..f4c42b2cdaddcc3da909f95c47a60edbcc10c764 --- /dev/null +++ b/build/torch213-cu126-aarch64-linux/metadata.json @@ -0,0 +1,37 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "W9KNHE3CYVULeJE4sqVc1w+BRUcBtlj2izQn0ug1y1w=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch213-cu126-aarch64-linux/metadata.json.sigstore b/build/torch213-cu126-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..f2105fafb01b14855f1fc088ffc57099c1b313b3 --- /dev/null +++ b/build/torch213-cu126-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrKgAwIBAgIUOXxpljZSmkWoe0MVPm8Y710cDCowCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTIyWhcNMjYxMDE5MTMxNTIyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEi2hRa1SHS8BCsh9Sc3g8p0r8SwKfRDjWLtykpFq1iEbgxe/07uGfJKbCvWGSgnPJ7p9Yp9R3ae9bTTjWGUVGdaOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVERDDgAABAMARzBFAiEAyFmKIGIbGbSqSYAj3+48W1P3awvDFi0ldvwSagLJ+RUCICi2DjP/a2mHkZnjPbkzUfwCZglQYH06HzJtjkknBGiGMAoGCCqGSM49BAMDA2kAMGYCMQCyJK6N3XSDX+2nVOyVMTxW+LUVjV4zTMQr1YK285NpE8lCyuaSNIz4CpsKW0S7wakCMQDYxu7oMDLR0kLXLFdvNDqJdWAtvV9PhD7OH8vZ+HGaDvdz4o0Ln9agge439BWPqFY="},"tlogEntries":[{"logIndex":"3321256527","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415122","inclusionPromise":{"signedEntryTimestamp":"MEQCIBh81LjKb/tsx0DwAem23ttmTmseKnZsFNJgF4llnLNbAiBJsyr1zPjNaBsEydraysvdlFJmyn6E8PwN9jmD3UMQ3A=="},"inclusionProof":{"logIndex":"3199352265","rootHash":"dxs26QSZXHwJrPOmwf+FtIppfn/i9+wkrbdxkXC016I=","treeSize":"3199352284","hashes":["1+XL8eM6nv2KKDEMMC01s1nSpamGkCVCsOn11pSujAU=","KGBv+V1LDmiimz8vy5Mj16jhWGRUk01e1YmWfVaKQRY=","CeVhGRI+XbdYp8oPUla9N+v2XoC5HUksofyGbJUAPe0=","X+a9ig4PFoSD4k9FmgCdsbAFXJSx9foxR/rWPZRoKeI=","yG8iuvhmuMFqk+WMi63n7xUEfj4tfJwzhVAxusmzeAI=","1glRWoKNhk9bPi2xmMzJ0xPkOhXpSA8DgV/6N/jw1EQ=","4RUXAOyFICltl19ZA7fgeHdclpt7OOf+akPOoMGV5IE=","zRydiKTQ2lwubTShI6BKXR2DLH/300r/vwdstdu9bFs=","pvSOMtCoDY7VfblJPv4FHsSUD1IfWr1b9YUvy3C4W64=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199352284\ndxs26QSZXHwJrPOmwf+FtIppfn/i9+wkrbdxkXC016I=\n\n— rekor.sigstore.dev wNI9ajBEAiBuSGkK5rNMhmzojBnQA6geUudo4er4wvNrxJ/MbzMTAAIgDNbhI33jqalAONjCG/GGbHB6XVV6GhQWyREvwDLordI=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI2OGQzZjk3MTUwOTk2ZmYwYWRhMWZjN2MyYTdjZWRmOTE0OTVjMDg0NWIwNTRhMWU2N2E1NjgyM2E3NjlmZjM3In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNNLzFEYmJUUVZ6K3B3ZEEzcEQrKyt2b3I1eW9WNkN6eUg4bkdBNWVaWG1RSWhBSTBvUGNYL0ZpbHU2MjU4b1U1RjM5MjUxTDE2eURGWWc2cm04ZjJDVGE5SSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KTFowRjNTVUpCWjBsVlQxaDRjR3hxV2xOdGExZHZaVEJOVmxCdE9GazNNVEJqUkVOdmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJKZVZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVsNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZwTW1oU1lURlRTRk00UWtOemFEbFRZek5uT0hBd2NqaFRkMHRtVWtScVYweDBlV3NLY0VaeE1XbEZZbWQ0WlM4d04zVkhaa3BMWWtOMlYwZFRaMjVRU2pkd09WbHdPVkl6WVdVNVlsUlVhbGRIVlZaSFpHRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWU1JFUm5RVUZDUVUxQlVucENSa0ZwUlVGNVJtMUxTVWRKWWtkaVUzRlRXVUZxTXlzME9GY3hVRE5oZDNaRVJta3diQXBrZG5kVFlXZE1TaXRTVlVOSlEya3lSR3BRTDJFeWJVaHJXbTVxVUdKcmVsVm1kME5hWjJ4UldVZ3dOa2g2U25ScWEydHVRa2RwUjAxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKclFVMUhXVU5OVVVONVNrczJUak5ZVTBSWUt6SnVWazk1VmsxVWVGY3JURlZXYWxZMGVsUk5VWEl4V1VzeU9EVk9jRVU0YkVNS2VYVmhVMDVKZWpSRGNITkxWekJUTjNkaGEwTk5VVVJaZUhVM2IwMUVURkl3YTB4WVRFWmtkazVFY1Vwa1YwRjBkbFk1VUdoRU4wOUlPSFphSzBoSFlRcEVkbVI2Tkc4d1RHNDVZV2RuWlRRek9VSlhVSEZHV1QwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgL8k4uNYAGP0BVAQOWCADR5DMSW39Si4YfUZPZatAuloCFH5MevmvAyihmqvmh+kLStXAVmNwGA8yMDI2MTAxOTEzMDUyMlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTMwNTIyWjAvBgkqhkiG9w0BCQQxIgQg0mjZ4sfYS8+B8KjSwowLRxwyTgc/6HTFWeJzHPghfq0wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxALBjnlbeJ2knEOAAtQZhbRUb62QTG5285eGYyJztupSLEa5/V23fW6b6VBbOxl5XFAIxAP90CuZmryrDT4J8l6NqvNroH0MYOqU43DCARGrNUnLL96jnjN2pGpCLRuIremTrlw=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"aNP5cVCZb/Ctofx8Knzt+RSVwIRbBUoeZ6VoI6dp/zc="},"signature":"MEYCIQCM/1DbbTQVz+pwdA3pD+++vor5yoV6CzyH8nGA5eZXmQIhAI0oPcX/Filu6258oU5F39251L16yDFYg6rm8f2CTa9I"}} \ No newline at end of file diff --git a/build/torch213-cu126-aarch64-linux/symbols.json b/build/torch213-cu126-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch213-cu126-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch213-cu130-aarch64-linux/__init__.py b/build/torch213-cu130-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch213-cu130-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch213-cu130-aarch64-linux/_ops.py b/build/torch213-cu130-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch213-cu130-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch213-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch213-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..1feaa68622e2c4dc75e0ebe241bb1356de68f376 --- /dev/null +++ b/build/torch213-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c1ed8261990ff7834dd43733e5652f462e00d17ca7750801ce14b5b2648213f3 +size 421408 diff --git a/build/torch213-cu130-aarch64-linux/metadata.json b/build/torch213-cu130-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..6c215c18a6d6c6e1c84ed218b75d9f7920648bc4 --- /dev/null +++ b/build/torch213-cu130-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "we2CYZkP94NN1Dcz5WUvRi4A0XyndQgBzhS1smSCE/M=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch213-cu130-aarch64-linux/metadata.json.sigstore b/build/torch213-cu130-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..5082954576db5effc592337f48b2f83d13c84c32 --- /dev/null +++ b/build/torch213-cu130-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrKgAwIBAgIUQKbMjQz/JGHVcrbt/11EmJRwmOEwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTIzWhcNMjYxMDE5MTMxNTIzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEyJvR5dLM5kAlKw0y3WE7I5Z3aU8rgUxYVTL2KUIM4FTVWWMfuPnEwRNfd0HYEHN0+/C7doZh3EjpreiuHFF/EqOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVERHCgAABAMARzBFAiA5g3vZj7nSwTjQIQq+H+0Q57llHr57dY7s50eq/Fg9BgIhAImoZ9VJUyr/NLDg3gZ76z2MJVPk+q0yg2NDaejphA/VMAoGCCqGSM49BAMDA2gAMGUCMAPYohBnzB1KFmPrcH/vevMG/H3D8LARYnpIjLRvtLAvr4HRZ+a0Sh92uc14HZdiIAIxAK+TQFVnZcwXoUvjkB/XMm5bv56KT2/e1nrUVyPqQ8sMkPWUlMXT3654ZsMhqpBVCg=="},"tlogEntries":[{"logIndex":"3321256676","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415123","inclusionPromise":{"signedEntryTimestamp":"MEQCIC2Oc+kGxmxrA0Dgbqz3vVBEoFtHTvGsjBfMr/qtVbv1AiBKR3y215qXVZ6P+2hLoli31cfcrDg6G/m37Zsa0lSJzg=="},"inclusionProof":{"logIndex":"3199352414","rootHash":"z298oPUrERwk4AtXRFOi+DVwKFav0yKyxwgQ84ue4xI=","treeSize":"3199352420","hashes":["//LeaFG4hATtFLjNcnT/VRquOqYl930pl/pyuuNgMs8=","IHlUjwn7xGCwWZoL5PMrEEsN8F3xoKdVIjGSvDidwhM=","X2kH+63hg3Jo3bfTbCxWjh8M42+InQJpp1knF/VCkR4=","JV4Vefo9UAwpR06ZiveNfrU32srTWIblOwP4U97DSWI=","MpAU1kXuoltMWZcBaSJXzXIxovuQ85sgGWrLEHVdEZg=","L06+0mxmtHeVwhmwb7t0uOfXa0nPB8lMBoNR1E+taMo=","8/hWmytFDrTzJcbLALfmAKUVivFdCbGvGvp0fe2+/UA=","AXCNG83+oxhlZbnRkiY5qDuMMQAyraexb7fxH+uAINQ=","pvSOMtCoDY7VfblJPv4FHsSUD1IfWr1b9YUvy3C4W64=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199352420\nz298oPUrERwk4AtXRFOi+DVwKFav0yKyxwgQ84ue4xI=\n\n— rekor.sigstore.dev wNI9ajBEAiBwqxLpMEYDqplHUnnbM2JIVTeqCGqiH6lNukrjOMydFQIge2KSkshdnDT9p6qUzCe48gaahVnumOzM/kDz+1hh9J4=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjMjc3YTA2MmNkNTExZDA3MzlhODQ4MjE1NzAxNDI5ZGRlOWFjODRlNzRjNmE1ODhjNTk2NDdjOThlMTcxMzkzIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNPcG01YVRQd2wrWGVqellrOUdDSzBHUUt1cjIrd2tJZmRHcUw3OUF2d0JnSWhBSW5wSkFkMVVBdUZlcG5TS25OTW4rbmx6bEV4eDRZTzE5UmcrZGVaM3VJNSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KTFowRjNTVUpCWjBsVlVVdGlUV3BSZWk5S1IwaFdZM0ppZEM4eE1VVnRTbEozYlU5RmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJKZWxkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVsNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVY1U25aU05XUk1UVFZyUVd4TGR6QjVNMWRGTjBrMVdqTmhWVGh5WjFWNFdWWlVURElLUzFWSlRUUkdWRlpYVjAxbWRWQnVSWGRTVG1aa01FaFpSVWhPTUNzdlF6ZGtiMXBvTTBWcWNISmxhWFZJUmtZdlJYRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWU1NFTm5RVUZDUVUxQlVucENSa0ZwUVRWbk0zWmFhamR1VTNkVWFsRkpVWEVyU0Nzd1VUVTNiR3hJY2pVM1pGazNjd28xTUdWeEwwWm5PVUpuU1doQlNXMXZXamxXU2xWNWNpOU9URVJuTTJkYU56WjZNazFLVmxCckszRXdlV2N5VGtSaFpXcHdhRUV2VmsxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKblFVMUhWVU5OUVZCWmIyaENibnBDTVV0R2JWQnlZMGd2ZG1WMlRVY3ZTRE5FT0V4QlVsbHVjRWxxVEZKMmRFeEJkbkkwU0ZJS1dpdGhNRk5vT1RKMVl6RTBTRnBrYVVsQlNYaEJTeXRVVVVaV2JscGpkMWh2VlhacWEwSXZXRTF0TldKMk5UWkxWREl2WlRGdWNsVldlVkJ4VVRoelRRcHJVRmRWYkUxWVZETTJOVFJhYzAxb2NYQkNWa05uUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgPsd3LlbMeqBztuHvukHiBoZT9VFY383PfBqu9UauoEMCFQCc+VlM/tUKMzpotd683Ou0oXhfAhgPMjAyNjEwMTkxMzA1MjNaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTEzMDUyM1owLwYJKoZIhvcNAQkEMSIEIDjwdtm2Q3HkA7F6F/TNeKzFLpBf1O9BKM7bS2DjRpzeMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMDnu6QbsGfUHbgO1EsaBSuwoHzbtHL0HxH41zj4+cPAZrWS4dfonH0KeRSP8V+nduQIwJ+sW28Wvg+xFYqa39MHTAL+9mu6CGOQdNOTYkwYp8wApJcBZY2GlevwTw/UDg+b1"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"wnegYs1RHQc5qEghVwFCnd6ayE50xqWIxZZHyY4XE5M="},"signature":"MEYCIQCOpm5aTPwl+XejzYk9GCK0GQKur2+wkIfdGqL79AvwBgIhAInpJAd1UAuFepnSKnNMn+nlzlExx4YO19Rg+deZ3uI5"}} \ No newline at end of file diff --git a/build/torch213-cu130-aarch64-linux/symbols.json b/build/torch213-cu130-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch213-cu130-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch213-cu132-aarch64-linux/__init__.py b/build/torch213-cu132-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch213-cu132-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch213-cu132-aarch64-linux/_ops.py b/build/torch213-cu132-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch213-cu132-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch213-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch213-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..a251e5693fdb64ecd04f9c452ded91a62149b327 --- /dev/null +++ b/build/torch213-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:34fbf3bf81f525803cd6d0993b7dc6e045468c6d9ba516177658d151aa34d854 +size 489256 diff --git a/build/torch213-cu132-aarch64-linux/metadata.json b/build/torch213-cu132-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..c4d25ec8327960e8082ea30ae5c86b4f38bc8335 --- /dev/null +++ b/build/torch213-cu132-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "NPvzv4H1JYA81tCZO33G4EVGjG2bpRYXdljRUao02FQ=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch213-cu132-aarch64-linux/metadata.json.sigstore b/build/torch213-cu132-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..9c176944b5079929693105cdb823baaf44c4f221 --- /dev/null +++ b/build/torch213-cu132-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrOgAwIBAgIUO+ku0RT0/3Zf5nVy/xTBCJqzXzswCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTI0WhcNMjYxMDE5MTMxNTI0WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEV9qh+k0tdSjTb/qI6FytrSZb23uxyyypqBG+Oe9LWdVi5VkZt19TGiPzXkX0WRx/oGV76RpV6tIEgGY3FdQKoqOCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVERK2wAABAMASDBGAiEAhu6GT/p86a6bSV2r1I/SNqtQPKOK0qHG+9FfC8fEGnwCIQCzsVaDPP9cQEidclmgFPsMJ3hXbwRcepXlRlaniScIzDAKBggqhkjOPQQDAwNoADBlAjBZlAP/D+dmGZmvLs5yjNodeUJ8sNxL9X0Mn/dhB7+sflnUBOzwq9iCnlH+9e1yiDcCMQCyk3FvuPw/UY7gCbi1f6u/MtYC+IM4/gP+e8GPTxqHHdc4xG+9MFfbLzOz+HwBx8A="},"tlogEntries":[{"logIndex":"3321256829","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415124","inclusionPromise":{"signedEntryTimestamp":"MEUCIBXE0N07VKCDnHdbEJkFEYHIMHpG42yC8k/dh2l2PtLpAiEAu+TU2Tby+MedO5Jg8msOs3Vk+oQJT4cBbNFodd5AWZs="},"inclusionProof":{"logIndex":"3199352567","rootHash":"yy8iMwcocGlylnan6kttlgwQ6GqMtSwNtZtFLa27C2I=","treeSize":"3199352677","hashes":["SsVg09RXc5KYFGH1MdFJowgYaRvBUcVBeKd6RTDmKjM=","tIOS1yrFJXWxq1ZAKc5bqDapBOaRS/B5T7Z+QozN5y0=","90mMxuWcNFE55FoCpP2TPrJc4wf3zNHcHskU9HDDInI=","HcNsTNmSAy4/M656o92dMVvcPG8uNUa/G/lkw1f9zrA=","338rpGFe4vjlOVkvCkkvP1hk5hp/VTVkyRDSqq4K3aA=","j8Gw64P4gBSoL+pna4kJHXS9sB+dnwQiMJBA6PIjgbk=","xwK9F0agWzydmmcWYCV5x/IHBFCyWZoARrH0mUkdErE=","LT0J8a8RmgrDxTpLDq+I2hE5QvMFLLCq7/C1r0bmk9Q=","cOEXMWpHT5zIMYc2nEFVsBzYX6tKRZIiguRjgaK69nc=","AXCNG83+oxhlZbnRkiY5qDuMMQAyraexb7fxH+uAINQ=","pvSOMtCoDY7VfblJPv4FHsSUD1IfWr1b9YUvy3C4W64=","cP4l3rlp30O2J4tK9B/xoAeqH/RdNZIPK//vSax8mag=","2kBL89oKHsEw0T2bX5j7Z/fslRfQBgZ+oj5rWgze1x4=","QyrNHyflFz1E0Eym1jXWQcmF271vpJhslMNdsmF41dM=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199352677\nyy8iMwcocGlylnan6kttlgwQ6GqMtSwNtZtFLa27C2I=\n\n— rekor.sigstore.dev wNI9ajBFAiEAw3oVd9yLKKJlyGvCD3rZM+TC1CCkZO0gxIsEeYX7VrgCIAW5XJlgM8TKVBEuRhOykv8gckYTvNaiIssiFo12+YlE\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJlYjc0NzcxOGYzYzdmODQ1N2E4OWMzZjI0NDgyMDJkMjU5YzdlZTU5ZjhiNTY0ZWU2ODJhNzcyYTM0ZmI5NTg1In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJQS83VlBob05mY0lxZDBVUU1oc2tRTFB2UGt4RkRJUnpDZU9NVVh5eGFjWkFpQURRWTZLTTF1RjdsS3MrYVZJYWFrdHl5bnJBVTkxdW1SVTI5UFdHUnVGQVE9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KUFowRjNTVUpCWjBsVlR5dHJkVEJTVkRBdk0xcG1OVzVXZVM5NFZFSkRTbkY2V0hwemQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJKTUZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVrd1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZXT1hGb0syc3dkR1JUYWxSaUwzRkpOa1o1ZEhKVFdtSXlNM1Y0ZVhsNWNIRkNSeXNLVDJVNVRGZGtWbWsxVm10YWRERTVWRWRwVUhwWWExZ3dWMUo0TDI5SFZqYzJVbkJXTm5SSlJXZEhXVE5HWkZGTGIzRlBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWU1N6SjNRVUZDUVUxQlUwUkNSMEZwUlVGb2RUWkhWQzl3T0RaaE5tSlRWakp5TVVrdlUwNXhkRkZRUzA5TE1IRklSd29yT1VabVF6aG1SVWR1ZDBOSlVVTjZjMVpoUkZCUU9XTlJSV2xrWTJ4dFowWlFjMDFLTTJoWVluZFNZMlZ3V0d4U2JHRnVhVk5qU1hwRVFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2IwRkVRbXhCYWtKYWJFRlFMMFFyWkcxSFdtMTJUSE0xZVdwT2IyUmxWVW80YzA1NFREbFlNRTF1TDJSb1FqY3JjMlpzYmxVS1FrOTZkM0U1YVVOdWJFZ3JPV1V4ZVdsRVkwTk5VVU41YXpOR2RuVlFkeTlWV1RkblEySnBNV1kyZFM5TmRGbERLMGxOTkM5blVDdGxPRWRRVkhoeFNBcElaR00wZUVjck9VMUdabUpNZWs5NkswaDNRbmc0UVQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgaeRNKkT4uz/oJgCAMV6eTzkka6bmbpiVjQxeBSDy6FgCFQCP1aDBQQK9LM6Nw4VpClQIn58vaRgPMjAyNjEwMTkxMzA1MjRaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTEzMDUyNFowLwYJKoZIhvcNAQkEMSIEIPHPBEMi7e787KIClF7H3sbigzrDW9AbhG92fJq5fgL4MIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQCYBmsCZfM+o8QtP8KgmFZQaUFwfKEgjR+LA+d7mR5brWCKeL1gbFEAz5JKljQLJncCMQDj1qxEGj/cruiTvmr46STCodEdbOYZYMoU26nuenkkmUIWJEyck26kclps4v0C0nk="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"63R3GPPH+EV6icPyRIIC0lnH7ln4tWTuaCp3KjT7lYU="},"signature":"MEQCIA/7VPhoNfcIqd0UQMhskQLPvPkxFDIRzCeOMUXyxacZAiADQY6KM1uF7lKs+aVIaaktyynrAU91umRU29PWGRuFAQ=="}} \ No newline at end of file diff --git a/build/torch213-cu132-aarch64-linux/symbols.json b/build/torch213-cu132-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch213-cu132-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch214-cu126-aarch64-linux/__init__.py b/build/torch214-cu126-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch214-cu126-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch214-cu126-aarch64-linux/_ops.py b/build/torch214-cu126-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch214-cu126-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch214-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch214-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..3591def6ea38b3dd079dafcb18e8520d8f1cc2f2 --- /dev/null +++ b/build/torch214-cu126-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:868f4a703cc95a3739d8efe2ebffce1b41460d91feeb62af10e9da0fb5c0f632 +size 285952 diff --git a/build/torch214-cu126-aarch64-linux/metadata.json b/build/torch214-cu126-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..68759b255baabdb237fe63344eef633a9a8210d7 --- /dev/null +++ b/build/torch214-cu126-aarch64-linux/metadata.json @@ -0,0 +1,37 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "ho9KcDzJWjc52O/i6//OG0FGDZH+62KvEOnaD7XA9jI=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch214-cu126-aarch64-linux/metadata.json.sigstore b/build/torch214-cu126-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..a220aada13d5d3255e032fc7ca15d6ba75ebae46 --- /dev/null +++ b/build/torch214-cu126-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrKgAwIBAgIUfl1LGbCqeRjFNzu0Gqy+uzEDR34wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTI2WhcNMjYxMDE5MTMxNTI2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEEWJLSZW8woWO9qdb0nEiMpbZ9FL02oYX61bL10THO/b8JTlbMjOdSze5SmgIdHew3WqKBYVw92EgJyTw0Rh+GKOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVERSugAABAMARzBFAiAJMiRyp3LuIJm/1pkuRfRv361/Smm59gwN6Jx1E/CBUAIhALeGDLv6jhp5prEyeqgddqL3pwf5E9Q59sIR93ckYm7RMAoGCCqGSM49BAMDA2gAMGUCMGByJPmp8kzJLN7k1s2PXZFPBFKZQxRuExIHKpMm1y7QJtFElkGKFfdfwYKMwcR2bQIxANwFkQvcmrW4aNpWQqgQxCUwnJ+kJP60FB+CafTVZKRZBR7S5FqoDsgaqz0GCKRhdA=="},"tlogEntries":[{"logIndex":"3321257158","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415126","inclusionPromise":{"signedEntryTimestamp":"MEYCIQCc0lG3L0Az7teDrgpIgrIoOgnRMDS3r/75AhxRlrICUgIhAIQ0UXaGlZPijrGP3Ph/+frnDQ2Fc9ojDVkAekvrd/a1"},"inclusionProof":{"logIndex":"3199352896","rootHash":"DRCItcpIKb6pPwyZrpj9JaBGDGYsM2KaRP0b/FNyl6g=","treeSize":"3199353146","hashes":["BgJ8vkwPWH2E2ulKRtNsqmpcj332qqh8+rpi+IHgFcw=","seCB7xB21VJpGESutp863ywGt0ODRImvevLpWc1fbcQ=","QZMMGIq4uakgu2yWSFkS7Apic5vUi9mKF110ku5TbGw=","WuOahZtICp2YO+vKFcIpz5gkKvD4LiTSPGyTJMlkV0E=","j1ZTSfkpVJPMOsTFtavfTRlHSRmhJDCNexaTBkLA3PE=","D/piLEoQ1wha/5lELOFTsk+dHsd24N83DXKZ1MXdHs4=","wEnUhujkXcZWOcf0AgRZ0p/9JXFg0uG7pNim+QbyN20=","lEa0edkRyyAbjfj9M42VNeMm44SW0sHFD4wmjhDGQ+k=","cn91+XQ3MB+IjsxJUK4kuXIyPmbisBSZKxO2U7I1qIY=","D4DLcaqe2LA8J89GPpQCsbQWeMEQ/7BVV1fbvw9gGcQ=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199353146\nDRCItcpIKb6pPwyZrpj9JaBGDGYsM2KaRP0b/FNyl6g=\n\n— rekor.sigstore.dev wNI9ajBFAiAxJKCqqYfqFcFmIRuDF3zx7r3OVatUgSpUJDEs8gyy9gIhAOyPVT/K4ReKX5XgEbiBi5wp/gbNDOdBLo6NF8Uxv/33\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIyOTgyODI4MzM1MTJkNWVjYjZiZjJiNTRmMTE0MmM4NTk0OGYwMjg4MDgzNDgyZjBiMzc3ZDQ0MjBiMTcxODE1In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUN1TllIbXBUNzllc0hva3lYSmVndkdYUU1Ud09pZms4V3VBT09WYWNEaE1nSWhBSXVCeWFRa0I2dkVmSmVEdlJ3S3ZjbE53ZFhoVE1ITUlYUDdKOVpBbGxwOCIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KTFowRjNTVUpCWjBsVlptd3hURWRpUTNGbFVtcEdUbnAxTUVkeGVTdDFla1ZFVWpNMGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJKTWxkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVreVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZGVjBwTVUxcFhPSGR2VjA4NWNXUmlNRzVGYVUxd1lsbzVSa3d3TW05WldEWXhZa3dLTVRCVVNFOHZZamhLVkd4aVRXcFBaRk42WlRWVGJXZEpaRWhsZHpOWGNVdENXVlozT1RKRlowcDVWSGN3VW1nclIwdFBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWU1UzVm5RVUZDUVUxQlVucENSa0ZwUVVwTmFWSjVjRE5NZFVsS2JTOHhjR3QxVW1aU2RqTTJNUzlUYlcwMU9XZDNUZ28yU25neFJTOURRbFZCU1doQlRHVkhSRXgyTm1wb2NEVndja1Y1WlhGblpHUnhURE53ZDJZMVJUbFJOVGx6U1ZJNU0yTnJXVzAzVWsxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKblFVMUhWVU5OUjBKNVNsQnRjRGhyZWtwTVRqZHJNWE15VUZoYVJsQkNSa3RhVVhoU2RVVjRTVWhMY0UxdE1YazNVVXAwUmtVS2JHdEhTMFptWkdaM1dVdE5kMk5TTW1KUlNYaEJUbmRHYTFGMlkyMXlWelJoVG5CWFVYRm5VWGhEVlhkdVNpdHJTbEEyTUVaQ0swTmhabFJXV2t0U1dncENVamRUTlVaeGIwUnpaMkZ4ZWpCSFEwdFNhR1JCUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgjM1WWeWWRBfJrq7YHOEta0c1MowuDI9Xv8PiBueqbQ4CFAUDkyHz9a5u7GeMfoCzpuPOyqsqGA8yMDI2MTAxOTEzMDUyNlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTMwNTI2WjAvBgkqhkiG9w0BCQQxIgQgK87qf9o8xzZeizEe8ipJMWyf3tgQBQ2CbwIz++w/yMEwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAPkN0JwzPhvDiD7ARWP/7W4DCfNPHAQvnNeUvkhoWyJjaTs4mf/aX6qJNoyoNGDDWgIxAKb8c+Jib62tBLx+A0khPWeYNbRz3YD4QvuTWq95UXw4iMaaplXoOxUn3V1GYJUZPw=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"KYKCgzUS1ey2vytU8RQshZSPAogINILws3fUQgsXGBU="},"signature":"MEYCIQCuNYHmpT79esHokyXJegvGXQMTwOifk8WuAOOVacDhMgIhAIuByaQkB6vEfJeDvRwKvclNwdXhTMHMIXP7J9ZAllp8"}} \ No newline at end of file diff --git a/build/torch214-cu126-aarch64-linux/symbols.json b/build/torch214-cu126-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch214-cu126-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch214-cu130-aarch64-linux/__init__.py b/build/torch214-cu130-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch214-cu130-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch214-cu130-aarch64-linux/_ops.py b/build/torch214-cu130-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch214-cu130-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch214-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch214-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..6ad7fc686660b837d64793fb3b8b5732c18ee1e2 --- /dev/null +++ b/build/torch214-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6b395947d1dedb45f048a3a46921e006dc7c4f81cb9f1fa6a956198d2f3cf0eb +size 421272 diff --git a/build/torch214-cu130-aarch64-linux/metadata.json b/build/torch214-cu130-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..a275bf5e565e65f9ca52ba47c885ec683fcc8838 --- /dev/null +++ b/build/torch214-cu130-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "azlZR9He20XwSKOkaSHgBtx8T4HLnx+mqVYZjS888Os=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch214-cu130-aarch64-linux/metadata.json.sigstore b/build/torch214-cu130-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..14e7c4384124227cf35fc67d203098b2e6aa9e49 --- /dev/null +++ b/build/torch214-cu130-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrKgAwIBAgIUC8EAT9QqmMfXbyhx4OvLHfdxab0wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTI4WhcNMjYxMDE5MTMxNTI4WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAENhxJ389IiRoo0c7m24YyYSRtfATcMW8c/5SITYSvSrwZwQDngK/RtwLKgnCVuXv0GSideNV05AI5nyTUWeibXqOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVERdSQAABAMARzBFAiEAhQVQDnSDeXCE4C9WpcCZ8tk7FbWpTLq21u3Xw5z5Aa8CIC7YAy45oHuJ0Fb+wyuiJSwQpK6PJdNJyDfypMlIiCrWMAoGCCqGSM49BAMDA2kAMGYCMQC6XiSZPuiHoY1kelnShUlpKNw63mRv5te/t1U3ACE6vm55qxYysb/rjlSSGeyWyPwCMQDe5FgL1gpe2se24bBR4BM/ARzepie/fBJejDEDzc3Vn/8NyOU+tdWEDwhCoLabzR4="},"tlogEntries":[{"logIndex":"3321257715","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415129","inclusionPromise":{"signedEntryTimestamp":"MEQCIFB7UeTXy/7cLu5xq6PfjrYYrSV6Wx+uVHnc7P6JQaJzAiB/kadfznS5dIlyVNex8yuVR0mjIT8RC/KdIwXqgtnnMQ=="},"inclusionProof":{"logIndex":"3199353453","rootHash":"xQ3J3fq7QI6C6g8P0JeMyrSBbwoD4L0wJRGsNghxo7Q=","treeSize":"3199353681","hashes":["vYpoQ79mHBqBgaA0/M7gN40aqPkUblzE70zZ65ryuD8=","WtpqzjNVkTp/ZLJ0xvjgijpBCQXLbhpmcISd/eDlYfE=","FXV0q7EqjVUHgUMDMdXJDsmQ22hyQmh9WV3Yzz35dbQ=","7MGiYVzalKjnSAImsygh81nD8DzXR/zgicQ1cGDVibo=","o1Tda55kVU5WBJKtyf1wr/6Im9MTIub/tOYMuI/wcqY=","NOnh93L+W02bX2E4QgTUHDYOsKW23ukaE0vu3l0zKRo=","pl1x9mGIg1KvC/Dqb4WzHHVNCGUktfl1G51zYnHmp1I=","6Vk+uUyYtmrvBMEuOqIT0HaUSR+rTzl+VAtCkNY7aVI=","8hrmRCN/Ud4ECOwXXSBsJ9Ul8Zu54WwibRXGTTOKXyY=","1/uAspg0C7JJhuHL2oD7yBkqYGA/kTgSEgJzrbOPF8Q=","D4DLcaqe2LA8J89GPpQCsbQWeMEQ/7BVV1fbvw9gGcQ=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199353681\nxQ3J3fq7QI6C6g8P0JeMyrSBbwoD4L0wJRGsNghxo7Q=\n\n— rekor.sigstore.dev wNI9ajBGAiEA0Hs8NngI/78Kp9y7N/aJtIByLxy2q+IRZTcnYJOM3WACIQCNUG6J9DHGbqB4VRGmvnJ8FMymWGLvb2CxyRA0lkGS5g==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJiY2ZhMzczYWJmMGZiMjU4OGZlYjI5NzM5N2UwZDg4YWFkMGY0ZDc3ZGIwZDQ5ODViMzIyNzZlOTRkOTYwNGFmIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUURGZzFHYjB3UHNrMVpBanJTSXVRNlVuSHdsMXJjanUvQ1FxZjc1YWhETEtnSWdQb3JHeVJubFNtNzJEMVFFYXh2ZEt4dDNCODYweEhCT2JDMU5WNVIzWW93PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KTFowRjNTVUpCWjBsVlF6aEZRVlE1VVhGdFRXWllZbmxvZURSUGRreElabVI0WVdJd2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJKTkZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkVrMFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZPYUhoS016ZzVTV2xTYjI4d1l6ZHRNalJaZVZsVFVuUm1RVlJqVFZjNFl5ODFVMGtLVkZsVGRsTnlkMXAzVVVSdVowc3ZVblIzVEV0bmJrTldkVmgyTUVkVGFXUmxUbFl3TlVGSk5XNTVWRlZYWldsaVdIRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWU1pGTlJRVUZDUVUxQlVucENSa0ZwUlVGb1VWWlJSRzVUUkdWWVEwVTBRemxYY0dORFdqaDBhemRHWWxkd1ZFeHhNZ294ZFROWWR6VjZOVUZoT0VOSlF6ZFpRWGswTlc5SWRVb3dSbUlyZDNsMWFVcFRkMUZ3U3paUVNtUk9TbmxFWm5sd1RXeEphVU55VjAxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKclFVMUhXVU5OVVVNMldHbFRXbEIxYVVodldURnJaV3h1VTJoVmJIQkxUbmMyTTIxU2RqVjBaUzkwTVZVelFVTkZOblp0TlRVS2NYaFplWE5pTDNKcWJGTlRSMlY1VjNsUWQwTk5VVVJsTlVablRERm5jR1V5YzJVeU5HSkNValJDVFM5QlVucGxjR2xsTDJaQ1NtVnFSRVZFZW1NelZncHVMemhPZVU5VkszUmtWMFZFZDJoRGIweGhZbnBTTkQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgipknUi3pVFmy08nFNM3FX4fcAYnB+1C0WL1ZCv8mIrICFQCf2ucemFL06WfcxIZwo3OOdyGO7hgPMjAyNjEwMTkxMzA1MjhaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTEzMDUyOFowLwYJKoZIhvcNAQkEMSIEILFaLolvZKGp5x1zp5V443LhfrPOYnJPmz1SgJYfR6rjMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMG00+7XqoqKHQS7uajVABYx/KvUpq5HxfXdtcGDUZDpn48wQ481PsT0KqO9vHOjHrwIxAM6l5SvV5/ut0aCYHzkhgcsApsLy0bEkIl9Is5UCYByDtCuF+raPpHOgTstpLSvSrQ=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"vPo3Or8PsliP6ylzl+DYiq0PTXfbDUmFsyJ26U2WBK8="},"signature":"MEUCIQDFg1Gb0wPsk1ZAjrSIuQ6UnHwl1rcju/CQqf75ahDLKgIgPorGyRnlSm72D1QEaxvdKxt3B860xHBObC1NV5R3Yow="}} \ No newline at end of file diff --git a/build/torch214-cu130-aarch64-linux/symbols.json b/build/torch214-cu130-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch214-cu130-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch214-cu132-aarch64-linux/__init__.py b/build/torch214-cu132-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch214-cu132-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch214-cu132-aarch64-linux/_ops.py b/build/torch214-cu132-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch214-cu132-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch214-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch214-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..8d9d26a697081fe74cf9e8979f92d5ee618f5d75 --- /dev/null +++ b/build/torch214-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:48e3754833211dee822aebe947e28c348a319095a4c363157696a645a1ad8cad +size 489120 diff --git a/build/torch214-cu132-aarch64-linux/metadata.json b/build/torch214-cu132-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..98cf13a85448b249b83b178fab8c803b8dfaf0ad --- /dev/null +++ b/build/torch214-cu132-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "SON1SDMhHe6CKuvpR+KMNIoxkJWkw2MVdpamRaGtjK0=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch214-cu132-aarch64-linux/metadata.json.sigstore b/build/torch214-cu132-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..8592a7c73bd0b419ae8a5439484a60446d6fe4ee --- /dev/null +++ b/build/torch214-cu132-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHKjCCBrGgAwIBAgIUcgpYIApBJ9DFqq2lmRfZvGc7SGwwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTMxWhcNMjYxMDE5MTMxNTMxWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEvOy7BSANkyjwL02GiooKtR/0k3H5C14y3d5szIl5XmpXyvdD/e+lySBDu1+W+G1Ogk+yV0i9vX/9SHOQ+UWCdqOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVERm9QAABAMARjBEAiBAcByJrJ4qWfcgwDrdiR40P3jZTZx+CX9Lr/HbMe7jKgIgGiIQ1jWIZKa5stHfyQi1PiM/au2DrPI0GYyY4K3Iu4YwCgYIKoZIzj0EAwMDZwAwZAIwDrGVLuyELu03CzSo5m+RKKn0SS1HxDNuogO5FBf5P01Vx23oL3tY8HEIc1oXqYQqAjAy1H6Z9giAVUl9geC7GXfMJKJhjf/CQrTI+nv2X+/MYmHHta3tnffi/pq73dDE/cs="},"tlogEntries":[{"logIndex":"3321258169","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415131","inclusionPromise":{"signedEntryTimestamp":"MEUCIF75/nQCw1SHlrG2RU7a+uNnjpaiqGVBF9JswhXeV9diAiEArhDar9SMf/5YP0z8XVyYcsw/YXllfHFXpbkPeIJIcis="},"inclusionProof":{"logIndex":"3199353907","rootHash":"tX3KzYif0V3Oc3rCfznhrEp7TaQVngr7DUj8NZAPm30=","treeSize":"3199353974","hashes":["BuDVod8FsvOmWLlwuqbHAbdDIGtlEPjcXabsPq/gJzs=","q0AsRKGQ2kSriIhT1bed3+/ssWWWXBJNMduKav9w9O8=","EjE6LoEHTcR4oNJPt03pN6ch7QmGiRcOCAhROK2qcYc=","PovAO0NvqkXBs5JIBoovfta/v6wJ/Kiox7JxP2maF5o=","cDLrpxKBT1siYv9Dkxnnh0smuJlZ+/BuYXT6HCN11pg=","4nKzvikgfHl94Ctp4ZrLr9czBIfaH4ZyRsJ+YI49GEs=","PSAm8wjtHwa05umG/O4TocVl0iWEmKXy5psJ4yFNlD4=","eXzF+SDXPV+MhidDCO64SQg2cmbkXD+XbyW2RLCEMm0=","D4DLcaqe2LA8J89GPpQCsbQWeMEQ/7BVV1fbvw9gGcQ=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199353974\ntX3KzYif0V3Oc3rCfznhrEp7TaQVngr7DUj8NZAPm30=\n\n— rekor.sigstore.dev wNI9ajBFAiAhnj1GaR6k2I3x/0yDapnVRo2WU5cJEXTxpnY8UYYfUAIhAPtRBv1CjyqPQZN4OG2ChzbwSy2HskXymOEEirv0Ru3T\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIwYTk5MTFjYzc1YjU4N2IzYjMzYWVhOTU0ZDgxY2QyNjY4NTQyYTViMjU3MTgzODIwNWNhNmZhNDZhZWExNDhlIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUURQU2NLZ0FkM3dsTDJwZkl2UkdXTWt3MWJNQng0RlllUFV5QldLK2dDMk1RSWdibnNBSytFS3o2WDVLTENnLzNobjFpWDFuNDB3T2hpeEZHS1ROcUtmVVdNPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExha05EUW5KSFowRjNTVUpCWjBsVlkyZHdXVWxCY0VKS09VUkdjWEV5YkcxU1pscDJSMk0zVTBkM2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJOZUZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkUxNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVYyVDNrM1FsTkJUbXQ1YW5kTU1ESkhhVzl2UzNSU0x6QnJNMGcxUXpFMGVUTmtOWE1LZWtsc05WaHRjRmg1ZG1SRUwyVXJiSGxUUWtSMU1TdFhLMGN4VDJkckszbFdNR2s1ZGxndk9WTklUMUVyVlZkRFpIRlBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWU2JUbFJRVUZDUVUxQlVtcENSVUZwUWtGalFubEtja28wY1ZkbVkyZDNSSEprYVZJME1GQXphbHBVV25nclExZzVUQXB5TDBoaVRXVTNha3RuU1dkSGFVbFJNV3BYU1ZwTFlUVnpkRWhtZVZGcE1WQnBUUzloZFRKRWNsQkpNRWRaZVZrMFN6TkpkVFJaZDBObldVbExiMXBKQ25wcU1FVkJkMDFFV25kQmQxcEJTWGRFY2tkV1RIVjVSVXgxTURORGVsTnZOVzByVWt0TGJqQlRVekZJZUVST2RXOW5UelZHUW1ZMVVEQXhWbmd5TTI4S1RETjBXVGhJUlVsak1XOVljVmxSY1VGcVFYa3hTRFphT1dkcFFWWlZiRGxuWlVNM1IxaG1UVXBMU21ocVppOURVWEpVU1N0dWRqSllLeTlOV1cxSVNBcDBZVE4wYm1abWFTOXdjVGN6WkVSRkwyTnpQUW90TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQge6Q5RR9HRf3HjfjI8iVl9hXcRnfFFgG3CyU67nD57uwCFQDPxxyq6l4ALQzK5iF41i/nUUsU0xgPMjAyNjEwMTkxMzA1MzFaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTEzMDUzMVowLwYJKoZIhvcNAQkEMSIEIHzKcgz7CkTa0n4ZnMWN4wbNrkSJAxrZaYokF4lOBNAbMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMB3dpDKG9IwVYjSkOVcyOlcB3jOFqPxNhQ5NpLtUacaGiH5JwBzmkIXOMkUg4qVh/gIwHIO7RJI1ayY5MFZGnRFwBx2i30gYqO0HG+glDlxtOiFQwkMAzZu1O78H3FYif2JM"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"CpkRzHW1h7OzOuqVTYHNJmhUKlslcYOCBcpvpGrqFI4="},"signature":"MEUCIQDPScKgAd3wlL2pfIvRGWMkw1bMBx4FYePUyBWK+gC2MQIgbnsAK+EKz6X5KLCg/3hn1iX1n40wOhixFGKTNqKfUWM="}} \ No newline at end of file diff --git a/build/torch214-cu132-aarch64-linux/symbols.json b/build/torch214-cu132-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch214-cu132-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch215-cu130-aarch64-linux/__init__.py b/build/torch215-cu130-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch215-cu130-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch215-cu130-aarch64-linux/_ops.py b/build/torch215-cu130-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch215-cu130-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch215-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch215-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..b776c8b22455abb332ec9b48bce5cf7b5579b6fe --- /dev/null +++ b/build/torch215-cu130-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:eb431fbd51d584312e4c27b3d8919bb97225d4cb306f41c1b939dbc08acaf826 +size 421200 diff --git a/build/torch215-cu130-aarch64-linux/metadata.json b/build/torch215-cu130-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..4d5eefb78c0b383a9e309621c25e608aa2386768 --- /dev/null +++ b/build/torch215-cu130-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "60MfvVHVhDEuTCez2JGbuXIl1Mswb0HBuTnbwIrK+CY=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch215-cu130-aarch64-linux/metadata.json.sigstore b/build/torch215-cu130-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..4d991199dd62c87df5ab66627087b9f970229640 --- /dev/null +++ b/build/torch215-cu130-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrOgAwIBAgIUfKJ6fvOqJLB2fMw7dyItCf1mRSYwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTMyWhcNMjYxMDE5MTMxNTMyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEBqvvDO7ANwa/iO+YSNzYQV0oXMVhcUQg/T+nOgFFT4B8gob3DnHNXzVmPkbHON5XyqD/Vd1PvnQOz1Uff4f/qqOCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVERsYAAABAMASDBGAiEAila0CZk/fV8Jxmn965FE9Oj9DdOrEQeLReY5gF9f/t0CIQCw6/TvkKCGWPH2ChgMQiEjFCIRypaD61qtAqjiXnlKdjAKBggqhkjOPQQDAwNnADBkAjAMzwcvNDU45IjRmcki9Zctp+orAlwzEBdOfg/JZIzCoicdbz2Ws6J3KoP/CKj9afwCMA0bdjUjPe/M2niWC2GstQZnMdOjrLKKA82vOINIdjNSIbCJE7FU0of0b25oZSbA4A=="},"tlogEntries":[{"logIndex":"3321258354","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415133","inclusionPromise":{"signedEntryTimestamp":"MEUCIChamdqeIlcC7g33shCpmDLdFr5CoKxjjqH5cfgazYVJAiEAhBZdiSpOoK8yI5kVY5/tJJqrAych8uelesxIhUotkeA="},"inclusionProof":{"logIndex":"3199354092","rootHash":"U9Vo258jXSmT2e34OtWal5FFFCGJYodiwL7pxAeJZHs=","treeSize":"3199354263","hashes":["RNC4jonLK8SEfVMnO2v7b4vUidggX3NfIJwdneGIjP0=","VenjiHoiBQVS23jZXQM728BigM9EjayEzV7SiJO8eqM=","Cq+fUyvQYxIZ+o1X/bDUAy8Z/KxzxqnJVtrWyAlL6TI=","laeghLDmzKHwe3bPLU+QgiDyKIH9ozotcKYyLb505xM=","h3SwMGGxGYZQ9jb3bQHa65Co0Hi/0pwN+evjcyiXQbs=","GZ6EdgPQrqvrCfg1V6YiODjLR5ZKnzMXyojZJkjOc3c=","H1XyaxEeNcczL/YJRqQZW2yqHh5FkX0GMNzzJYj7/O4=","rBYVTUL3LkllX34on3aQKzSFp3NrxTN2X6OWnw/JL8o=","NnkiK9Sn4TCNrzEADH+HJVeTl/c97/fO7KBBK754Wvs=","eXzF+SDXPV+MhidDCO64SQg2cmbkXD+XbyW2RLCEMm0=","D4DLcaqe2LA8J89GPpQCsbQWeMEQ/7BVV1fbvw9gGcQ=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199354263\nU9Vo258jXSmT2e34OtWal5FFFCGJYodiwL7pxAeJZHs=\n\n— rekor.sigstore.dev wNI9ajBFAiBjZ6fdi+OZcJGFuC1AUfPfcP3UIEn3RnHBGVgJ8PMx8wIhAJv3c1lHRwzHObKl/yC/RUuAv83p6CsnLMBIYg5bqt2T\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI0OTM0ZTM0MThhZDJlMzViMTgyMTNkMjZiYzFmMTQwOTVkM2ZhMGU5ZTgzZWZkNGM0MTc0NmQzNTg5NDA5YzYyIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNTNVB2ZnFBK3I2OUV3Q2tTc08rS3JFOStoNU56QnhEd2dTMzJWcHhQc09nSWdPOXNUOGJsLzJJbmE0NUdSbDlKdW95T1hBUnBGZ01FQmVMTVVBWGdaSEtJPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KUFowRjNTVUpCWjBsVlprdEtObVoyVDNGS1RFSXlaazEzTjJSNVNYUkRaakZ0VWxOWmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJOZVZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkUxNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZDY1haMlJFODNRVTUzWVM5cFR5dFpVMDU2V1ZGV01HOVlUVlpvWTFWUlp5OVVLMjRLVDJkR1JsUTBRamhuYjJJelJHNUlUbGg2Vm0xUWEySklUMDQxV0hseFJDOVdaREZRZG01UlQzb3hWV1ptTkdZdmNYRlBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWU2MxbEJRVUZDUVUxQlUwUkNSMEZwUlVGcGJHRXdRMXByTDJaV09FcDRiVzQ1TmpWR1JUbFBhamxFWkU5eVJWRmxUQXBTWlZrMVowWTVaaTkwTUVOSlVVTjNOaTlVZG10TFEwZFhVRWd5UTJoblRWRnBSV3BHUTBsU2VYQmhSRFl4Y1hSQmNXcHBXRzVzUzJScVFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2JrRkVRbXRCYWtGTmVuZGpkazVFVlRRMVNXcFNiV05yYVRsYVkzUndLMjl5UVd4M2VrVkNaRTltWnk5S1drbDZRMjlwWTJRS1lub3lWM00yU2pOTGIxQXZRMHRxT1dGbWQwTk5RVEJpWkdwVmFsQmxMMDB5Ym1sWFF6SkhjM1JSV201TlpFOXFja3hMUzBFNE1uWlBTVTVKWkdwT1V3cEpZa05LUlRkR1ZUQnZaakJpTWpWdldsTmlRVFJCUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQguqbzKcRi6ocJgK+tfof0vCLzwvAfTkxjebq6NORcBzUCFBOT7UnoNUw3h8y/WttucDTv18coGA8yMDI2MTAxOTEzMDUzMlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTMwNTMyWjAvBgkqhkiG9w0BCQQxIgQgK8YwwYqbrIv0G8M/DXfqeT2OVxVFf8sdGNtfkfJ6g+YwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAKn2XfklDWuAR/UlbZCyQnYNaXV5x0ckZCFxQmDOag36cCUAJUr+pp+FpzkbxZwm9AIxANupjMYNr75YgJmnG7rvdOtMSyPyXnzEMCKOBlxj6BZ6Uye17oeY8a126+cz6RI0Eg=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"STTjQYrS41sYIT0mvB8UCV0/oOnoPv1MQXRtNYlAnGI="},"signature":"MEUCIQCS5PvfqA+r69EwCkSsO+KrE9+h5NzBxDwgS32VpxPsOgIgO9sT8bl/2Ina45GRl9JuoyOXARpFgMEBeLMUAXgZHKI="}} \ No newline at end of file diff --git a/build/torch215-cu130-aarch64-linux/symbols.json b/build/torch215-cu130-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch215-cu130-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch215-cu132-aarch64-linux/__init__.py b/build/torch215-cu132-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch215-cu132-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch215-cu132-aarch64-linux/_ops.py b/build/torch215-cu132-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch215-cu132-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch215-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch215-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..d1d39f19d1905650c4cfc8a09e3eaa13831b1e50 --- /dev/null +++ b/build/torch215-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:98ca4a62191b38821f47ec505fbf3d9e1e8a8d48f980b12bfd9b35ac97fe06f1 +size 489048 diff --git a/build/torch215-cu132-aarch64-linux/metadata.json b/build/torch215-cu132-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..70da15f99fe1e95abd024f115c6f283697f67b41 --- /dev/null +++ b/build/torch215-cu132-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "mMpKYhkbOIIfR+xQX789nh6KjUj5gLEr/Zs1rJf+BvE=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch215-cu132-aarch64-linux/metadata.json.sigstore b/build/torch215-cu132-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..47d2a235a45d7cd961c8a1ead5807213e1faf65b --- /dev/null +++ b/build/torch215-cu132-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrGgAwIBAgIUWUwH87Mcj5A5LSi1HeNMWzEh3LQwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTM1WhcNMjYxMDE5MTMxNTM1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEhQU4Z1yTl59SKR9su0rxZDKnub5Gmino5W1yTO92IFLR+vQdA3VNYnkiLY3A5Xjhkm/r1KkrIVZoL3bRgubJ4qOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVER2pgAABAMARjBEAiAHZlUOutrU4gys5GzMV3xHnKgcp2PJpz+kddvxSk646wIgcTULSTdnmFAja8Ikw2cID2wJyllN0yPPINIueU21ipswCgYIKoZIzj0EAwMDaQAwZgIxAI4yscXKcajQNTX0vWLUxiIQtZm9Z5eSyhXAL38jWUsdwbWrTWCJPBum5kgY6MXDfAIxAMbo8yuWZu9HMTLayjQu9Ux5Kuc7JIXsxCrRXFwVGzUbT6k6S/uK+SmFBpxgEGvd7w=="},"tlogEntries":[{"logIndex":"3321258927","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415136","inclusionPromise":{"signedEntryTimestamp":"MEQCIBoOZSP8vuBT7MJlganjqmdfBTnT8WVx/jnSNdf43n2TAiBsD0LahY7tLJpI/i4+ss+HThzajJ+XZGHohwrWnPkxpA=="},"inclusionProof":{"logIndex":"3199354665","rootHash":"81rE3fRKxB/iEfTzoNLqRaWCvxVLNzBIk1ejtg/C6ww=","treeSize":"3199354869","hashes":["2cROAAxdj9kJI4eOXf5ccMloMlr4jAb1gpUMLqi7+fQ=","49IHpyKcRwPwcVNeIcAYdWPEwczGdyBFrBPmSvKpSjU=","9jRfYfp73bhahSCn15Yx+fVtNiuqwbBNUvNjmfU9m3Y=","wbq720ZUImEFpGK3mVWORT+9cWR/Z4sM9nxYY5R3Kr8=","GPYoF36ZS/DAMPN6MpgJYEL625Wl880O6W1sJiJftzo=","bsJxQ+689opNpJq2ZMSYshmO4rG4JdO1WSCvFRhLwQU=","FwXrda8QrVdDqtBenwEYzAsP7txRATpH0jVeQjBuoRQ=","4iGCHj9r1ntRvNMxHzPex/HljlHIg6GVS0AqWQnyS2c=","HOoCExcHbNuz1vVR+OkcnBqTxTGXGCMJiq4I2Lq4TpE=","zMVGYOwkGlnSXpWYmuaUWoZ0nP0J9HfllAEoMNyJ4rI=","eXzF+SDXPV+MhidDCO64SQg2cmbkXD+XbyW2RLCEMm0=","D4DLcaqe2LA8J89GPpQCsbQWeMEQ/7BVV1fbvw9gGcQ=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199354869\n81rE3fRKxB/iEfTzoNLqRaWCvxVLNzBIk1ejtg/C6ww=\n\n— rekor.sigstore.dev wNI9ajBFAiAX3Mv3bX5XZeAdmRg82wVtZ2Lkof8ljvbEDwzn1u++nAIhAID6/Xqf1iUyLIBXDY4dv7Un8iNHbC3aQv0jBrhXJJpN\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI5YWFjMzdjYzVhM2NmMzVhZTk5Zjk2Y2Q1ZDNmZTJkNzhlMjQ3YmVlYWYxOTM0MTlhYWQzMjMyOTdmZDI3MzEwIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNBYXpFeEZGdFdKTXdta1l4b0M2QmRHRVBEQ05VSFQxMHNqeUIyTmJHdW1RSWhBTE5Nblk5WGhHSmJ1Z1IyVFZia01hdUpwY2FNZHVndnpoUWFvMVVOTnpwYiIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KSFowRjNTVUpCWjBsVlYxVjNTRGczVFdOcU5VRTFURk5wTVVobFRrMVhla1ZvTTB4UmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJOTVZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkUweFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZvVVZVMFdqRjVWR3cxT1ZOTFVqbHpkVEJ5ZUZwRVMyNTFZalZIYldsdWJ6VlhNWGtLVkU4NU1rbEdURklyZGxGa1FUTldUbGx1YTJsTVdUTkJOVmhxYUd0dEwzSXhTMnR5U1ZaYWIwd3pZbEpuZFdKS05IRlBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWU01uQm5RVUZDUVUxQlVtcENSVUZwUVVoYWJGVlBkWFJ5VlRSbmVYTTFSM3BOVmpONFNHNUxaMk53TWxCS2NIb3Jhd3BrWkhaNFUyczJORFozU1dkalZGVk1VMVJrYm0xR1FXcGhPRWxyZHpKalNVUXlkMHA1Ykd4T01IbFFVRWxPU1hWbFZUSXhhWEJ6ZDBObldVbExiMXBKQ25wcU1FVkJkMDFFWVZGQmQxcG5TWGhCU1RSNWMyTllTMk5oYWxGT1ZGZ3dkbGRNVlhocFNWRjBXbTA1V2pWbFUzbG9XRUZNTXpocVYxVnpaSGRpVjNJS1ZGZERTbEJDZFcwMWEyZFpOazFZUkdaQlNYaEJUV0p2T0hsMVYxcDFPVWhOVkV4aGVXcFJkVGxWZURWTGRXTTNTa2xZYzNoRGNsSllSbmRXUjNwVllncFVObXMyVXk5MVN5dFRiVVpDY0hoblJVZDJaRGQzUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQggq1jFIlj+EbxINJOjfVDE7sxaz8m/2sJXLXl9oRX5MgCFG4pStzePtxT8QltJ5tautETke8ZGA8yMDI2MTAxOTEzMDUzNVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTMwNTM1WjAvBgkqhkiG9w0BCQQxIgQglbqIfZDMny6YDnn6nbHc2NImgMM1UuW+wLYD+dtTLLYwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAODs2cH+p5asXF9gl2QC1f8btox0bq61KqRh4G/aBIOjCS8CNWrGgb/qY+iNVaYxvwIxALb5pdIF/HBRIMbNdS1p8RQnlLU3skgIDFF/dyNyxfmdDqXivU+aSx2PUXWB72GeXg=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"mqw3zFo881rpn5bNXT/i144ke+6vGTQZqtMjKX/ScxA="},"signature":"MEYCIQCAazExFFtWJMwmkYxoC6BdGEPDCNUHT10sjyB2NbGumQIhALNMnY9XhGJbugR2TVbkMauJpcaMdugvzhQao1UNNzpb"}} \ No newline at end of file diff --git a/build/torch215-cu132-aarch64-linux/symbols.json b/build/torch215-cu132-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch215-cu132-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +} diff --git a/build/torch215-cu134-aarch64-linux/__init__.py b/build/torch215-cu134-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch215-cu134-aarch64-linux/__init__.py @@ -0,0 +1,170 @@ +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/torch215-cu134-aarch64-linux/_ops.py b/build/torch215-cu134-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch215-cu134-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +import torch +from . import _rwkv_cuda_fe2bb66 +ops = torch.ops._rwkv_cuda_fe2bb66 + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"_rwkv_cuda_fe2bb66::{op_name}" diff --git a/build/torch215-cu134-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch215-cu134-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..d77904f319b87d65282f5cd6c448f6cb5d33a6c8 --- /dev/null +++ b/build/torch215-cu134-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c9debc7387eb11fc964c68094bba0cf578bac7529e4febc34703c162fc8b3739 +size 488088 diff --git a/build/torch215-cu134-aarch64-linux/metadata.json b/build/torch215-cu134-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..019fbcc260e96fc28bc1b69dbcf0b9b378a46a53 --- /dev/null +++ b/build/torch215-cu134-aarch64-linux/metadata.json @@ -0,0 +1,39 @@ +{ + "name": "rwkv", + "id": "_rwkv_cuda_fe2bb66", + "version": 1, + "kernels-minver": "0.14.0", + "license": "Apache-2.0", + "python-depends": [], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0", + "8.0", + "8.9", + "9.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "__init__.py": "amGCZ7NM7A7uSNkWvcloOWf4aT83qrs8SqDXoTS6QKQ=", + "_ops.py": "w0bKbLzY87IOlz0PXuprnKLp2kh3iAzHURNEjrzNcSI=", + "_rwkv_cuda_fe2bb66.abi3.so": "yd68c4frEfyWTGgJS7oM9Xi6x1KeT+vDRwPBYvyLNzk=", + "symbols.json": "inVs3pRoEtQOuFmFelx4SLt4ypDMYcHgUZsPgauo4+M=" + } + }, + "provenance": { + "kernel-builder": { + "version": "0.18.0-dev0", + "commit": "d054906a13710b811b96697b4635f03020842733", + "dirty": false + }, + "kernel": { + "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "dirty": false + } + } +} \ No newline at end of file diff --git a/build/torch215-cu134-aarch64-linux/metadata.json.sigstore b/build/torch215-cu134-aarch64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..37f67f621a721696b499f77c9cb8ed36f958e298 --- /dev/null +++ b/build/torch215-cu134-aarch64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHKjCCBrGgAwIBAgIUY2+aC/4ZqTbKav1v8QJeEsL0/WEwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTMwNTM4WhcNMjYxMDE5MTMxNTM4WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEnbCQcDs/YWI/J+ZjIRZN0rHCovS1Y6bL2OK+O1E489j74MMoce6ORv5zKZZqoLcpB4c0VpR70YmH/lqYR/MLQKOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVESDQAAABAMARjBEAiAhcdj7ncxD5GSmjEpur1FDzb0JIEAySmzYbNqSGVDwKgIgSP6vqqERd9nuhM1bygm13EZDhW3j470AbQWOn9RMdnswCgYIKoZIzj0EAwMDZwAwZAIwQme1XcUhEH6/wgYRbaXy0YbpJexEnd7taBqbt8/BhDalrNvr2WFgZxtAnN2/exmCAjAjIusKSdlqwkVkDOAR1EQsnTF79GCLZb1213q05rJDMdDgo4r2uqAJajeFxdt/y0o="},"tlogEntries":[{"logIndex":"3321259384","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792415138","inclusionPromise":{"signedEntryTimestamp":"MEUCIHp2hAuaQ05q5ZJ/bb0l23sxj1rOZvplVfowBEKrTLGTAiEA4LxIc1SLp8+As9lwhXfEDywE48KnvudadLYg+eW025o="},"inclusionProof":{"logIndex":"3199355122","rootHash":"DS2BgU5X8EHk8+UKC7ofS8QBkZ45WV8fd+IH8YTX+ms=","treeSize":"3199355376","hashes":["fh8nrGQ6dXXRFjUahDsirkqOmqnXNqiYuTw9dF0iO2Y=","C+0rbMGU3NsR518cmBMj533LLwhLFikDDTEHGq4rLGE=","hYiUpM4NwPlE6jt8zWdk/bY7bP/6g0QttPEcEIlipw4=","9SWRzOwST5heDCWNA2qR8Grtt+wn4gEV2GSJUVySmhU=","oh9z3dx9JWRkKSq95GQDAh8vjyeWhcNWOpzKP3C/m1s=","rObynn0awUbiviI/oyyEbi+1xmlGC23mg4Y5NB9cz98=","IBikIeiwF2VGijSgbrZdYGJoO/ezTAV2rjS+4ra3AWg=","VsHVnsBE83LsVJ7rZ0IOovyjf0h0+N8GsNXflF+/ZFU=","W54tpW1z6q8x0+nlP1D6e3rYhyyVPh2cOvTU11v7feg=","zctqo3RNjMFOG2qL1kgFaFDlgaB9KvsBfbQy3/wxcf4=","D4DLcaqe2LA8J89GPpQCsbQWeMEQ/7BVV1fbvw9gGcQ=","DlkFsuaW0Lemnx/hfWKDDGIHO8n35fi5A3sL1nZM6Nw=","kHa2Ke03zePMosAfwPNg2y4kk1ZwjdZi4MKNVX7epiE=","uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=","nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3199355376\nDS2BgU5X8EHk8+UKC7ofS8QBkZ45WV8fd+IH8YTX+ms=\n\n— rekor.sigstore.dev wNI9ajBGAiEA79eUhHVD7BX3G3ZZQR1i/LJ5RJY3ukgfJbv1m15uWJcCIQDAFgs/e1N8YRo/iDqqEHGpXtxa076P9lcHDeoPMtyZmA==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3ZDQyOTA2Yzc2N2UxYTJjYWU2NjM0Y2UzYjNiZTY2YjIyOGMyY2IyNzUwZDExN2I5NDQzYWYzZDY2MzE0YTBjIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRUFpbjdwYStIUzZZOHV1L0dQVVh5cExRRnAwWEZ2TytqVFRnbVFFTmJsNUFpRUFxTnVKYTB3eEU0a2VJSXB2blFKeEZ3VUlGcDIxOFV5b0MzZHFRNTNJT3F3PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExha05EUW5KSFowRjNTVUpCWjBsVldUSXJZVU12TkZweFZHSkxZWFl4ZGpoUlNtVkZjMHd3TDFkRmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFMTNUbFJOTkZkb1kwNU5hbGw0VFVSRk5VMVVUWGhPVkUwMFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ1WWtOUlkwUnpMMWxYU1M5S0sxcHFTVkphVGpCeVNFTnZkbE14V1RaaVRESlBTeXNLVHpGRk5EZzVhamMwVFUxdlkyVTJUMUoyTlhwTFdscHhiMHhqY0VJMFl6QldjRkkzTUZsdFNDOXNjVmxTTDAxTVVVdFBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmtWVFJGRkJRVUZDUVUxQlVtcENSVUZwUVdoalpHbzNibU40UkRWSFUyMXFSWEIxY2pGR1JIcGlNRXBKUlVGNVUyMTZXUXBpVG5GVFIxWkVkMHRuU1dkVFVEWjJjWEZGVW1RNWJuVm9UVEZpZVdkdE1UTkZXa1JvVnpOcU5EY3dRV0pSVjA5dU9WSk5aRzV6ZDBObldVbExiMXBKQ25wcU1FVkJkMDFFV25kQmQxcEJTWGRSYldVeFdHTlZhRVZJTmk5M1oxbFNZbUZZZVRCWlluQktaWGhGYm1RM2RHRkNjV0owT0M5Q2FFUmhiSEpPZG5JS01sZEdaMXA0ZEVGdVRqSXZaWGh0UTBGcVFXcEpkWE5MVTJSc2NYZHJWbXRFVDBGU01VVlJjMjVVUmpjNVIwTk1XbUl4TWpFemNUQTFja3BFVFdSRVp3cHZOSEl5ZFhGQlNtRnFaVVo0WkhRdmVUQnZQUW90TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgUrMvCCgLvrWxpZb6BnPHxfqZA/hqqycVlaUtGEXKrucCFQDVVQU6yLZu3tLgTOAQCF3/tkn/+xgPMjAyNjEwMTkxMzA1MzhaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTEzMDUzOFowLwYJKoZIhvcNAQkEMSIEIEf31qzMCtnV8K7XvZBCLCqIvmLCmEGiXlu3Nv88HLfCMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMDKmyxArSc61BHK6nwI411fn+45blp1dvUDCpoi9S+TpVfhELBc89OPfQbiXQN8m1AIwJRipcbVsxvcLQmNi6PPad+7tDpbeX/pUbi/4ENLXiaKmJuQZ7nSZmBZp4Ijpl6CH"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"fUKQbHZ+GiyuZjTOOzvmayKMLLJ1DRF7lEOvPWYxSgw="},"signature":"MEUCIEAin7pa+HS6Y8uu/GPUXypLQFp0XFvO+jTTgmQENbl5AiEAqNuJa0wxE4keIIpvnQJxFwUIFp218UyoC3dqQ53IOqw="}} \ No newline at end of file diff --git a/build/torch215-cu134-aarch64-linux/symbols.json b/build/torch215-cu134-aarch64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch215-cu134-aarch64-linux/symbols.json @@ -0,0 +1,331 @@ +{ + "schema_version": 1, + "module": "_rwkv_cuda_fe2bb66", + "functions": [ + { + "name": "forward", + "qualname": "forward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (float32).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_bf16", + "qualname": "forward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass (bfloat16 inputs/outputs, float32 ``w``).\n\nRuns the CUDA kernel and writes the result into ``y`` in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state", + "qualname": "forward_with_state", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (float32).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.float32``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "forward_with_state_bf16", + "qualname": "forward_with_state_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV forward pass with persistent state (bfloat16 inputs/outputs, float32 ``w`` and ``s``).\n\nRuns the CUDA kernel using and updating state ``s`` and writes the result into ``y``.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u: Input tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n k: Key tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n v: Value tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n y: Output tensor, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n s: Stateful tensor, shape ``[B, C]``, dtype ``torch.float32`` (updated in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B`` and ``C``; ``y`` shares ``[B, T, C]`` with inputs.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "s", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward", + "qualname": "backward", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (float32).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.float32``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.float32``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.float32`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.float32`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + }, + { + "name": "backward_bf16", + "qualname": "backward_bf16", + "module": "_rwkv_cuda_fe2bb66", + "kind": "function", + "doc": "RWKV WKV backward pass (bfloat16 inputs/outputs/gradients, float32 ``w``).\n\nWrites gradients into the provided tensors in-place.\n\nArgs:\n w: Decay weights, shape ``[C]``, dtype ``torch.float32``.\n u, k, v, y: Forward-pass tensors, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gy: Gradient of ``y``, shape ``[B, T, C]``, dtype ``torch.bfloat16``.\n gw: Gradient for ``w``, shape ``[C]``, dtype ``torch.bfloat16`` (written in-place).\n gu, gk, gv: Gradients for ``u``, ``k``, ``v`` respectively, shape ``[B, T, C]``, dtype ``torch.bfloat16`` (written in-place).\n\nNotes:\n - All tensors must be on the same CUDA device.\n - Shapes must agree on ``B``, ``T`` and ``C``.", + "signature": { + "parameters": [ + { + "name": "w", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "u", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "k", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "v", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "y", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gy", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gw", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gu", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gk", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + }, + { + "name": "gv", + "kind": "POSITIONAL_OR_KEYWORD", + "default": null, + "annotation": "torch.Tensor" + } + ], + "return_annotation": "None" + } + } + ], + "layers": [] +}