diff --git a/.gitattributes b/.gitattributes index da3f0c04fb9096fcd38e3e313272f98982136af4..2b34ecb3fa2929422a702ed627002ab6488def50 100644 --- a/.gitattributes +++ b/.gitattributes @@ -270,3 +270,21 @@ build/torch214-cu132-aarch64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lf 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 +build/torch210-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch210-cu128-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch210-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch211-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch211-cu128-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch211-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch212-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch212-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch212-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch213-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch213-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch213-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch214-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch214-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch214-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch215-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch215-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text +build/torch215-cu134-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so filter=lfs diff=lfs merge=lfs -text diff --git a/build/torch210-cu126-x86_64-linux/__init__.py b/build/torch210-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch210-cu126-x86_64-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-x86_64-linux/_ops.py b/build/torch210-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch210-cu126-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch210-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..6459e99c4e8521a4cdcef9fdca9faf2259afe8e5 --- /dev/null +++ b/build/torch210-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6d79862b4955ad5f1324d92752a3cf73aa091d9e366b7dc492aff6da2936bd2a +size 229664 diff --git a/build/torch210-cu126-x86_64-linux/metadata.json b/build/torch210-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..44586225fa0b80cb59c27c0923a5b7d48cc6c43d --- /dev/null +++ b/build/torch210-cu126-x86_64-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": "bXmGK0lVrV8TJNknUqPPc6oJHZ42a33Ekq/22ik2vSo=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch210-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..179b4e03213fb6ffc1454cd8d5bd3d00f7f2ce77 --- /dev/null +++ b/build/torch210-cu126-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrKgAwIBAgIUfvmeRF5qaJL+qjlbjwyxmGR4w/kwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1ODQ1WhcNMjYxMDE5MTcwODQ1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEsWc/TOQKmrOqdNfAZg73CzCWuIXMdwz3QdBIXtdoTlmCS4ltx9g7JfD2d5qJnNE9e0jhfGRxEiVJsq2l872iTaOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRnvMwAABAMARzBFAiEAxg1Ev88DWwckchLb3VsjH+aNJy77ffEK1wLgDothj+sCIHptoVaDlwerBMFqfES3qWbpiQQMRv15a3sGUCIFWEkfMAoGCCqGSM49BAMDA2gAMGUCMQCwDNzflNdjkuNcR76WOqAF5+wavV/DfDDnQIyT2qV7ceWG9do2FUjkbExKPuTxiDUCMEZvSphwXAOx9QCc3NP5J78RQoscRWTuaz8IJL+mCY3Uu1PQX+04Z2AY34xOYwUD9w=="}, "tlogEntries":[{"logIndex":"3323770314", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429128", "inclusionPromise":{"signedEntryTimestamp":"MEUCIHp7TSjzrQhwngkVaNQmnjdupFFzLzvnadqPCKK+uN/OAiEA7HEK4LTmQ+rACXldE6LAo+fFdDqm6i0/dTMfas1s13A="}, "inclusionProof":{"logIndex":"3201866052", "rootHash":"q2LwnMa0qCwWmN9LeMzPr62STdmzlllQ6N91OE2W/fw=", "treeSize":"3201866384", "hashes":["5UDUgs3+k5Xzij1BfLtcEVp3m9ZDfV2deWq3Kb/LA7k=", "gdaHQGHZsuM7S2rXGjSGkZtE6+gahOoqf/skRr4H5do=", "5LK+HFnrU+YQnqRl2fNbwGuWl3Vja+R/mHdQpLrVqhE=", "m1CnsCnGvURDzd0Bok6nQuMGnSpklfxErViAwou5G3U=", "/fpN9Qt0y/qIaXCjoC3wUc4QITIcjGW+VBVEtT0ClK8=", "2vsL4Ebp1isI8u1im81yyE4PzbTjVxZwzvYQdfL15MU=", "XX2SUkUyLNGFLO2SdTrYlSpC3cnuS86YtOJbLglzS+w=", "oUFQDWnPO2mq03t2Ty0dl1YEZvggpbiul7i2g1gX8bw=", "/LHEawDTQYBFJbe9u8bKTgIBkdH++kRL37pjgiHAcSQ=", "CzgNb3o0du4vXIkkV2nIwgUUfnxrcRx2bmFpiQvRIA4=", "ZPauTSPOesHyNSsCSamwg8AVqbZFSiWGAQUOkU+jTk8=", "qn+fHiszAh+N2YtRmkvz0lgvQltc3SQw8JGks4iML24=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201866384\nq2LwnMa0qCwWmN9LeMzPr62STdmzlllQ6N91OE2W/fw=\n\n— rekor.sigstore.dev wNI9ajBFAiEAk7RXxOeNmTa1mvHHwq/p83Nbbn4Jp76TSFlPpqzrd/cCIHtbzUgSL3Fg7Rq+f7oFMvKnCSfU6+kXercduvH9cl9N\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI4OGQ2ODc4OWExYTFkOTVkMWNhN2FiMGNlYWJhN2M2NjQzMTk4ZjQzMTgxYzg3OTU2MzExZjE5ODgzZjExMWI2In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUURldTd5K0kyVytFenF0VjBtTWh2UnNBR3U0eVpRSkFraXV6VmRIVnlLTTlnSWhBT21mSDQ3NmRiSVRsVHQ5KzlIck9XcC96VENVQjNFcTBHOVZoc3R3OVJ2TyIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KTFowRjNTVUpCWjBsVlpuWnRaVkpHTlhGaFNrd3JjV3BzWW1wM2VYaHRSMUkwZHk5cmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMFJSTVZkb1kwNU5hbGw0VFVSRk5VMVVZM2RQUkZFeFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ6VjJNdlZFOVJTMjF5VDNGa1RtWkJXbWMzTTBONlExZDFTVmhOWkhkNk0xRmtRa2tLV0hSa2IxUnNiVU5UTkd4MGVEbG5OMHBtUkRKa05YRktiazVGT1dVd2FtaG1SMUo0UldsV1NuTnhNbXc0TnpKcFZHRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKdWRrMTNRVUZDUVUxQlVucENSa0ZwUlVGNFp6RkZkamc0UkZkM1kydGphRXhpTTFaemFrZ3JZVTVLZVRjM1ptWkZTd294ZDB4blJHOTBhR29yYzBOSlNIQjBiMVpoUkd4M1pYSkNUVVp4WmtWVE0zRlhZbkJwVVZGTlVuWXhOV0V6YzBkVlEwbEdWMFZyWmsxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKblFVMUhWVU5OVVVOM1JFNTZabXhPWkdwcmRVNWpVamMyVjA5eFFVWTFLM2RoZGxZdlJHWkVSRzVSU1hsVU1uRldOMk5sVjBjS09XUnZNa1pWYW10aVJYaExVSFZVZUdsRVZVTk5SVnAyVTNCb2QxaEJUM2c1VVVOak0wNVFOVW8zT0ZKUmIzTmpVbGRVZFdGNk9FbEtUQ3R0UTFrelZRcDFNVkJSV0Nzd05Gb3lRVmt6TkhoUFdYZFZSRGwzUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg8+fkk+U89KaCUe9mA1rORkwzULQ0BBSkS1n3SMWAqkACFQCEKcFlvapVRMb0vUbH6PM252PDchgPMjAyNjEwMTkxNjU4NDVaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE2NTg0NVowLwYJKoZIhvcNAQkEMSIEIKFN7I1rMTAxWElm2q6BbXaS+WgKV3wZb6x1aM70OCjiMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQD+1N/PkTeg5MBx25t7IWVyoGndI4hlUDsBGOhQQFi4Ios4iSwGSzVJJzNAOICb9O8CMQC+9/wB67JVKcMe9mSIuew4SFAhn8WjjCheVhvpGAMqKKp5CVEwFBTSfFyONlxQf2w="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"iNaHiaGh2V0cp6sM6rp8ZkMZj0MYHIeVYxHxmIPxEbY="}, "signature":"MEYCIQDeu7y+I2W+EzqtV0mMhvRsAGu4yZQJAkiuzVdHVyKM9gIhAOmfH476dbITlTt9+9HrOWp/zTCUB3Eq0G9Vhstw9RvO"}} \ No newline at end of file diff --git a/build/torch210-cu126-x86_64-linux/symbols.json b/build/torch210-cu126-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch210-cu126-x86_64-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-x86_64-linux/__init__.py b/build/torch210-cu128-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch210-cu128-x86_64-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-x86_64-linux/_ops.py b/build/torch210-cu128-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch210-cu128-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch210-cu128-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..a3604f62e864d3fe008e1a27627a7054aeb3e9e7 --- /dev/null +++ b/build/torch210-cu128-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b0da7af93dbb4efcad7603fe28da3ff189f4439d8733093c04c425c30d107657 +size 385160 diff --git a/build/torch210-cu128-x86_64-linux/metadata.json b/build/torch210-cu128-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..2c51c8bfb83276a6e0248b55dc7d31182de6ce52 --- /dev/null +++ b/build/torch210-cu128-x86_64-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": "sNp6+T27TvytdgP+KNo/8Yn0Q52HMwk8BMQlww0Qdlc=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch210-cu128-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..381ef838cca2469b3e7b2a257cdef056d9aedfa2 --- /dev/null +++ b/build/torch210-cu128-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLjCCBrOgAwIBAgIUEqiF9HCeB5HUNUIy3upxoOaeBgswCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1ODU0WhcNMjYxMDE5MTcwODU0WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEeiUlWC7nB/zjW06YXeGQwzz3tYRR2GtUC8ujyEIJSMUEUEzI3LfeKXvqYtQbDBL4HZjAmvM+qFNTG/IVhpOrn6OCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRoR9AAABAMASDBGAiEAz2gOMi2DFvTerYdCsDl1rbLizAgXLZl4ISNsWCLpjn8CIQDqg5U7JSBxBvCFEcyum3yn9nDd0Onnn1NFWkHOCVmjpzAKBggqhkjOPQQDAwNpADBmAjEAyq929zvmzo+pNRIVw5MpAeerUxyI99934SQxcrCoWYNbHtD3ooKJXWAT8l2l1Iv6AjEAraWCDw9XhhhEMsodWqXZd8avDfMJuawzwH3ndH/8/tvnAoauT8eVVNm6nQR/9w4I"}, "tlogEntries":[{"logIndex":"3323771599", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429134", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQDMfq4u7CdOqnUncJ4QvO9Cc2bF2sU/FNSXaLNIoZ1NNQIhAKfYIDmdPXcsCMzKNxrG9yfCAhT5rZ2zu8IYwPZuWQ3A"}, "inclusionProof":{"logIndex":"3201867337", "rootHash":"3fEzcA30cuBB2+hjnDLIygAZCttDjC4W1jKhaZv7xt8=", "treeSize":"3201867360", "hashes":["qmVuki3PmPGsTSS6i/KFmyiwjR8nw9cPAe28l7CZfYU=", "mBqP8Hl8GoYvYZhficTIvEns5o0/WUDhc6zOGbNtQzY=", "15QXL18gepKqb9eZvs1lh7yk9W6wxIu0TBmGz5nl6g4=", "/TU9mQZz/HjJpzzD+pfQZCYJ2rziEiwwbJoce9YfTag=", "S1biDtwzg5rTUMHIyaC8EU+sbwDW/n1aTwL+eONmRKI=", "5kETYltmJi2K3WWVc/CRxLs88C+UG1Hfg600H3vnhOA=", "8kaicoSUVjjSC/faCJu7SRyFm53BivN12STKOdBsHjc=", "WP8jqy3iiV+Qu6Dljlu5beMIdjdp6cfbVyjWOsZlxYM=", "ZPauTSPOesHyNSsCSamwg8AVqbZFSiWGAQUOkU+jTk8=", "qn+fHiszAh+N2YtRmkvz0lgvQltc3SQw8JGks4iML24=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201867360\n3fEzcA30cuBB2+hjnDLIygAZCttDjC4W1jKhaZv7xt8=\n\n— rekor.sigstore.dev wNI9ajBFAiEAlu6pd2BgxtQw/tVC9CCPBwVoB3bXRO9Sc5i3x4xXrusCIGCvo4pQ7JAYQQii0QWtTmy4fZ+BNZttGgkUYnvZi1UO\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjZDFmOGJmODhiNzYzYzZlMjk4ZDJkZjY5ZDdmZTZmMDNmN2I4MzFjNDI3OGU0ZmI0NWQxZDNjNGM0ZjA1NjM2In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRkErd1o5Sk9zeEhsOGpZbVI0NW5hUEVGdGdqRWxpU1VFQjRIV2JVY2ZyNkFpRUE5NDNOelRCRnN2R2VqNGtGZWZGOEc2eG1OKys5cTg1VnpiZGJlQWR6eVF3PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1ha05EUW5KUFowRjNTVUpCWjBsVlJYRnBSamxJUTJWQ05VaFZUbFZKZVROMWNIaHZUMkZsUW1kemQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMFJWTUZkb1kwNU5hbGw0VFVSRk5VMVVZM2RQUkZVd1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZsYVZWc1YwTTNia0l2ZW1wWE1EWlpXR1ZIVVhkNmVqTjBXVkpTTWtkMFZVTTRkV29LZVVWSlNsTk5WVVZWUlhwSk0weG1aVXRZZG5GWmRGRmlSRUpNTkVoYWFrRnRkazByY1VaT1ZFY3ZTVlpvY0U5eWJqWlBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKdlVqbEJRVUZDUVUxQlUwUkNSMEZwUlVGNk1tZFBUV2t5UkVaMlZHVnlXV1JEYzBSc01YSmlUR2w2UVdkWVRGcHNOQXBKVTA1elYwTk1jR3B1T0VOSlVVUnhaelZWTjBwVFFuaENka05HUldONWRXMHplVzQ1YmtSa01FOXVibTR4VGtaWGEwaFBRMVp0YW5CNlFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2NFRkVRbTFCYWtWQmVYRTVNamw2ZG0xNmJ5dHdUbEpKVm5jMVRYQkJaV1Z5VlhoNVNUazVPVE0wVTFGNFkzSkRiMWRaVG1JS1NIUkVNMjl2UzBwWVYwRlVPR3d5YkRGSmRqWkJha1ZCY21GWFEwUjNPVmhvYUdoRlRYTnZaRmR4V0Zwa09HRjJSR1pOU25WaGQzcDNTRE51WkVndk9Bb3ZkSFp1UVc5aGRWUTRaVlpXVG0wMmJsRlNMemwzTkVrS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg64SEkVed56Dbpk0+Yjj5pyF1Pn7icvrqlGElLCZHPOMCFQD9vndFTKBg/kYQZUh4BkfxL4W1zxgPMjAyNjEwMTkxNjU4NTRaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE2NTg1NFowLwYJKoZIhvcNAQkEMSIEIOFMJrmgSQiPX4CHIx8nhJXUiXajcnuTeHXHqQPuvAhZMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQDXULQ6IvTBBs85iyov2/SLC5+35Yc/5se24B2n5abySzm0HPPZrVYd4brJa+7ODdwCMQDv6fPUjw6LSvYvF8rvL/JpmtJBkhWpH1P/fYzIufzI9Oo6JG7X51D89zaafjAJbCA="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"zR+L+It2PG4pjS32nX/m8D97gxxCeOT7RdHTxMTwVjY="}, "signature":"MEUCIFA+wZ9JOsxHl8jYmR45naPEFtgjEliSUEB4HWbUcfr6AiEA943NzTBFsvGej4kFefF8G6xmN++9q85VzbdbeAdzyQw="}} \ No newline at end of file diff --git a/build/torch210-cu128-x86_64-linux/symbols.json b/build/torch210-cu128-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch210-cu128-x86_64-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-x86_64-linux/__init__.py b/build/torch210-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch210-cu130-x86_64-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-x86_64-linux/_ops.py b/build/torch210-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch210-cu130-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch210-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..a89cdb63aacc080afb2acc83f6325273c2dd91ae --- /dev/null +++ b/build/torch210-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:00e2ba59a3e2b328deb5523162a5bf961c8165f5df3c0edb994c88fef1740933 +size 410664 diff --git a/build/torch210-cu130-x86_64-linux/metadata.json b/build/torch210-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..bf74addf7ea4985eb7307012e73ac9683a09cae4 --- /dev/null +++ b/build/torch210-cu130-x86_64-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": "AOK6WaPisyjetVIxYqW/lhyBZfXfPA7bmUyI/vF0CTM=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch210-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..06ea12b3152a9db74e923bc37defd9e22c441ca4 --- /dev/null +++ b/build/torch210-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLjCCBrOgAwIBAgIUOsJwl36bHG0b5M6n6N8GIcDMOWMwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1ODU4WhcNMjYxMDE5MTcwODU4WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEOi7IC1tvZn3KGWj2DIyV+GfGYvIwLhE2InvV5AwvWCWHPQW2JmQUERMXsv2KlGqtwbwQ75/BdG//P9LQSYTlf6OCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRojqgAABAMASDBGAiEAx70P2tYJMB7zKUD4jlA9YpIK/a0US+pBdsn7qeSI/qICIQDQqSZisdWFlO5gHZYXzUNLNGQGsfcg3TKKK8d1ZPN3VjAKBggqhkjOPQQDAwNpADBmAjEAt3R3vygUrET/tv9JfQIV8/3Kzx6o7sK7UVkGPr6Qa3aNeQF5cMwHvasrQ/owUMbAAjEA+C0EV5iFuBQz9Nk0nwIZ1zQq9hN6/WIaD7JABM7A1BjzJDzJ+UxGW5CiOGZsPWEX"}, "tlogEntries":[{"logIndex":"3323772481", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429139", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQD2QYKDIKxOJeWibmgDyX/j02PrvbVKP2+H0dhrdJEkrAIgNIJrWdk2CdyaHfIxzIJ5rujxIms5HQ9n6fsQ9ntyYgY="}, "inclusionProof":{"logIndex":"3201868219", "rootHash":"qHbhsvcv93NY6aefgLbNIYE4kVUCUe0NV4eRVtccJys=", "treeSize":"3201868228", "hashes":["XAJBY/JKLpaY+dZxwc8PRU3xKJYDOa6jdgSHazEt58M=", "pNTlWqlrtrecvhv74ttcMhQMX0MqxU0fLEjQAWIx2ak=", "q7727o3jyfoDHP49If1sHmp5SKnGdIYcRIR4Ckyyxb4=", "VMnMz50AwbunonLX+4hOkb/BD7wMlUGq1EkqnayhtX4=", "+69j4S4C5nFafHUwgRbshWAqrthzGc/YRC8YVcH9Tmw=", "lg7U7emXAefNlKN53uhjb8I10Ih5efu+b+44izHUuWY=", "i2Y+y+mZHjhlULnWgcHFtScBO2yGOMG9AinRZJQXS1c=", "yKBUVDJgU8ACKPM6C5MHn6W30PlPzIXzq9ZBzDLHFR0=", "vIj1uM8VfWQMrUqihHhuL2Af8QkL7QJu14hc1Yg+dyc=", "KToNGOU/2LQPMnwkW0YBQmo27fLOVdrtlPNJA6eS4kg=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201868228\nqHbhsvcv93NY6aefgLbNIYE4kVUCUe0NV4eRVtccJys=\n\n— rekor.sigstore.dev wNI9ajBFAiAIMtQI6LfYFAs3TZVanZZLR0iR/Hw5G2jbJan6oK5exwIhAP0cG0sIHNan1ltP2wkIjG3FegVSJmkdCfM5Ey9+gzB/\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJhNTgxYThlMmEyZWQ5YzkzZTYzMGVjYzRlOTgyN2FlOGZjMTI0YjdjNGE4MGVjMzEyNWMwNDA5YjU4MDJlNWQ5In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJR0xxZXFSK2x6K2pMb0R3bW1HUEx3eWVJNWNYWmppNGM3cDNoNE5ndkc1WEFpRUF1OW5rMWRjQVdOZWtkTmFGSXpWUWszbEI3OEp1ZmlJeFFnUHYzZXRyQW9zPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1ha05EUW5KUFowRjNTVUpCWjBsVlQzTktkMnd6Tm1KSVJ6QmlOVTAyYmpaT09FZEpZMFJOVDFkTmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMFJWTkZkb1kwNU5hbGw0VFVSRk5VMVVZM2RQUkZVMFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZQYVRkSlF6RjBkbHB1TTB0SFYyb3lSRWw1Vml0SFprZFpka2wzVEdoRk1rbHVkbFlLTlVGM2RsZERWMGhRVVZjeVNtMVJWVVZTVFZoemRqSkxiRWR4ZEhkaWQxRTNOUzlDWkVjdkwxQTVURkZUV1ZSc1pqWlBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKdmFuRm5RVUZDUVUxQlUwUkNSMEZwUlVGNE56QlFNblJaU2sxQ04zcExWVVEwYW14Qk9WbHdTVXN2WVRCVlV5dHdRZ3BrYzI0M2NXVlRTUzl4U1VOSlVVUlJjVk5hYVhOa1YwWnNUelZuU0ZwWldIcFZUa3hPUjFGSGMyWmpaek5VUzB0TE9HUXhXbEJPTTFacVFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2NFRkVRbTFCYWtWQmRETlNNM1o1WjFWeVJWUXZkSFk1U21aUlNWWTRMek5MZW5nMmJ6ZHpTemRWVm10SFVISTJVV0V6WVU0S1pWRkdOV05OZDBoMllYTnlVUzl2ZDFWTllrRkJha1ZCSzBNd1JWWTFhVVoxUWxGNk9VNXJNRzUzU1ZveGVsRnhPV2hPTmk5WFNXRkVOMHBCUWswM1FRb3hRbXA2U2tSNlNpdFZlRWRYTlVOcFQwZGFjMUJYUlZnS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg4N4ZW2JIsnMQ8WktGQQfV54/81q5wmiFvn5xq1XqMYgCFQCUFAn+H0Eb7gKVwo8RiUboYZDSaBgPMjAyNjEwMTkxNjU4NThaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE2NTg1OFowLwYJKoZIhvcNAQkEMSIEIEUq50YKG1d1cOGD5WTXnP7L2vfL2/ynmuNt3yUm3nzsMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQCambyE1BpdvML0lsOd+a0lvvBM63fCWwD2M2+gmwbuTG1buGRFXIEIAdswRpCe1tACMQCZCwN0vcwByry50MSlPYR5iWWkI3Mobv4mWH4k5jWEaGnziKCJfU3z5wc1+MayDoo="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"pYGo4qLtnJPmMOzE6YJ66PwSS3xKgOwxJcBAm1gC5dk="}, "signature":"MEUCIGLqeqR+lz+jLoDwmmGPLwyeI5cXZji4c7p3h4NgvG5XAiEAu9nk1dcAWNekdNaFIzVQk3lB78JufiIxQgPv3etrAos="}} \ No newline at end of file diff --git a/build/torch210-cu130-x86_64-linux/symbols.json b/build/torch210-cu130-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch210-cu130-x86_64-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-x86_64-linux/__init__.py b/build/torch211-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch211-cu126-x86_64-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-x86_64-linux/_ops.py b/build/torch211-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch211-cu126-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch211-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..1a6ab691de8fcfbd2a1a4313df9f7472a775e177 --- /dev/null +++ b/build/torch211-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ab0e923dce74e7c55bb8abea8d9d8f7d96f186074471149103c1faefa1b417aa +size 229664 diff --git a/build/torch211-cu126-x86_64-linux/metadata.json b/build/torch211-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..79da078dab759b17ecc1a6bbec204c31118f0b68 --- /dev/null +++ b/build/torch211-cu126-x86_64-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": "qw6SPc5058VbuKvqjZ2PfZbxhgdEcRSRA8H676G0F6o=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch211-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..06ba9daa4de3bf01140413714d617316c7da7964 --- /dev/null +++ b/build/torch211-cu126-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrKgAwIBAgIUFVsVaQeTF27KyuMvRs303QB46QEwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1OTAzWhcNMjYxMDE5MTcwOTAzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEWqN3ZkXF3JPq6dO0re1kIMdlCGT+Byd8nirifO4PIcw/E+DmIPBuRjHUxv7ytgv54PO3nyygW6ivZoXrCpjbUqOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRo2vQAABAMARzBFAiA7KbDArvBqPe1ALnBUn5xPqTBrCu0uPZ5o5JIVNmtWWQIhAI72Pj0QRZdmAGaW9SPU+XY5a2TZJeFFz2TObTVLf+vsMAoGCCqGSM49BAMDA2kAMGYCMQCql53VlR9HerIP/8d8MS2TkDgps2e9xa9eubylKsvXZ3t+rOlrAzEWxd9VNcevsWcCMQChmdaNx1EI2HBcIBckusOpzZ3Qc5/Rrhvb7BaeWEFWRwhhuTgdNQ+GXa5mhEI+6fU="}, "tlogEntries":[{"logIndex":"3323773771", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429145", "inclusionPromise":{"signedEntryTimestamp":"MEQCID/VN2RjyUCVq0k+ZOsrG08jdEJR1aDaf4j3j1HyWIMkAiBCAPSRq+t7dlEc7EBUxkXOUL/QqCMVneeIp7AyUVljDw=="}, "inclusionProof":{"logIndex":"3201869509", "rootHash":"W9claMolZ2LXhc/23lO1qqtJ10niMaPMJZgoa7HKxEs=", "treeSize":"3201870071", "hashes":["raHjQY+QIo39yvdFKrg7KeRipu/iDdUoUxzfUW2d0/o=", "eRcYa8n7auHMoYlKergGcS7JuBFDy4PuFoVMCsLqgUg=", "/fD8z0DuRGkfJK9aH/KX7Npv+AeXsOTHDezmb9uBfw8=", "XPJy46K1U8x9i06aKdJZzOTFXAAiwqwCK4qCuJsvCew=", "d5Cw1FQX+FKrn7bb2yDe64FrJ8rI/fA7lVc1TrBIRxE=", "jqor0B/zlWyyxIp6l1vzgn7tdsR17YWEG7+SbewIwjM=", "JmlM7j6N/BrbY5MEwBtf4UBDf7ICI4oP8dQR3m8tPnQ=", "epPzYj28zI13Ex7N6T1AEFlj2e7OfhJNzUTft7YXXaw=", "GJdEUFVAjBjlgt2QrxlxbFPeFn+gGuCFQJC/BhEoVks=", "kAXcK3qZBVLwzxb8d5ZiNPgflqeJhoaRtsQea1HAhhc=", "/VTLdp40meMX8B3H3VAry72qvwJQ+KlULgQDrKWn4yU=", "sghJxkV4aXmhkDlPntrZ3JulU+RYeZIS3dl45uyZNe0=", "KToNGOU/2LQPMnwkW0YBQmo27fLOVdrtlPNJA6eS4kg=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201870071\nW9claMolZ2LXhc/23lO1qqtJ10niMaPMJZgoa7HKxEs=\n\n— rekor.sigstore.dev wNI9ajBFAiBaFRg68PGZFSz5CRU+Vj5ksZ9mA6Ydrr6ZUuvS+JHDpwIhAK0SfX8LL9Vfjwxndf7OewBCQDt1imaxHdWG+gDcMZao\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI1MDBlMDIxNjJlOTVlM2JiZmY3ODE5MzAxZmEwZWJhNzUzZmY2NmY4ODhmNjY5ZjcxY2ZmZGVkZDExNDY0YTE5In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJQWo5WVRBeXdRRTlCNXpYNzRyY2hEWU5BNDFUeFBWTExua1BJMDdFa3ZCWUFpQnBFNjBhR2VxaHI5WVJwZ3ptZmRDdnNNaVc3K2VDbGZYQkNZQ3BJdDZrQ0E9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KTFowRjNTVUpCWjBsVlJsWnpWbUZSWlZSR01qZExlWFZOZGxKek16QXpVVUkwTmxGRmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMVJCZWxkb1kwNU5hbGw0VFVSRk5VMVVZM2RQVkVGNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZYY1U0eldtdFlSak5LVUhFMlpFOHdjbVV4YTBsTlpHeERSMVFyUW5sa09HNXBjbWtLWms4MFVFbGpkeTlGSzBSdFNWQkNkVkpxU0ZWNGRqZDVkR2QyTlRSUVR6TnVlWGxuVnpacGRscHZXSEpEY0dwaVZYRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKdk1uWlJRVUZDUVUxQlVucENSa0ZwUVRkTFlrUkJjblpDY1ZCbE1VRk1ia0pWYmpWNFVIRlVRbkpEZFRCMVVGbzFid28xU2tsV1RtMTBWMWRSU1doQlNUY3lVR293VVZKYVpHMUJSMkZYT1ZOUVZTdFlXVFZoTWxSYVNtVkdSbm95VkU5aVZGWk1aaXQyYzAxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKclFVMUhXVU5OVVVOeGJEVXpWbXhTT1VobGNrbFFMemhrT0UxVE1sUnJSR2R3Y3pKbE9YaGhPV1YxWW5sc1MzTjJXRm96ZENzS2NrOXNja0Y2UlZkNFpEbFdUbU5sZG5OWFkwTk5VVU5vYldSaFRuZ3hSVWt5U0VKalNVSmphM1Z6VDNCNldqTlJZelV2VW5Kb2RtSTNRbUZsVjBWR1Z3cFNkMmhvZFZSblpFNVJLMGRZWVRWdGFFVkpLelptVlQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQghj+oPKYi7hKpdknu+CX5veQMys+Ly5bNqBV0dawdLR0CFQDMiWwTnoBddiRr0LzhpeiYiCqEYRgPMjAyNjEwMTkxNjU5MDNaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE2NTkwM1owLwYJKoZIhvcNAQkEMSIEIAtRNBFmvTLWUG/bF7MmAqALkL6Odc4YASmnl8iMgP0VMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMAapjIKaeDyH2pau/ieCmDLjNK7iKPcjTSj1SSH/oRaBO4Pq9jrdqaTLFMtOWNc/egIxAKr1/QNKkJ8gtomOyTsEtPwvLIW5vctBXkWj3oAhyUkHWOgvv81CFSkuxzzRcceSTw=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"UA4CFi6V47v/eBkwH6Drp1P/ZviI9mn3HP/e3RFGShk="}, "signature":"MEQCIAj9YTAywQE9B5zX74rchDYNA41TxPVLLnkPI07EkvBYAiBpE60aGeqhr9YRpgzmfdCvsMiW7+eClfXBCYCpIt6kCA=="}} \ No newline at end of file diff --git a/build/torch211-cu126-x86_64-linux/symbols.json b/build/torch211-cu126-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch211-cu126-x86_64-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-x86_64-linux/__init__.py b/build/torch211-cu128-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch211-cu128-x86_64-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-x86_64-linux/_ops.py b/build/torch211-cu128-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch211-cu128-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch211-cu128-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..301b82877b4049dd02f2b368901a22a92ba02ba5 --- /dev/null +++ b/build/torch211-cu128-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8d92b9e6c9a232d9dd98944890e941b4cfb40c2742be6026a3900c2a133567cd +size 385160 diff --git a/build/torch211-cu128-x86_64-linux/metadata.json b/build/torch211-cu128-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..e0f0e140c637140385c09c21ad4d50d97af7ee3e --- /dev/null +++ b/build/torch211-cu128-x86_64-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": "jZK55smiMtndmJRIkOlBtM+0DCdCvmAmo5AMKhM1Z80=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch211-cu128-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..87f9b4562563ea583f5e5cade3e0a4e55d16ee3f --- /dev/null +++ b/build/torch211-cu128-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrKgAwIBAgIUQsRIFi4slmKfyHUCXcP3W4xHCGEwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1OTEyWhcNMjYxMDE5MTcwOTEyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAElIOR5ypMJW6Ilbh/Z3CWfKPwrbIhF7gYD81zZGWe7tWbRrBIBZ7EC/QNFMtD3nXuUXZVZOiVzm78leOeRutKn6OCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRpaGQAABAMARzBFAiEAx3/bVETT68VXVa0guGum2aHhCpH0d48XyHfGRazWDlsCIAvXbFRkhoFMHHin56p4Dk/KwbiI3Ojz8ogfk64BotXGMAoGCCqGSM49BAMDA2kAMGYCMQC5dMLNK7wpSWILmStCeXl8dIox4UgGoz2GRXZfi1m7VlGn2DZaa7mYwaeatQ89IrICMQCzdBT9Mjye2FPtTpVsx036oiQxcZCM3qHfjyl0iBSo+m8wOnMyGLzeHSKtwToyB8s="}, "tlogEntries":[{"logIndex":"3323775407", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429153", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQCv9PEdQ2ohZRyhQdaZiGJiw+vHRm9rctQnfHKSik5dcgIgVx9HYlhg24aJD07+qWkiL6fPMRlZdWaO4SsEGxB2Bug="}, "inclusionProof":{"logIndex":"3201871145", "rootHash":"54qruKG0yzo+f75Kpm/kL2m289lYdeozLCDD9m9SNOQ=", "treeSize":"3201871354", "hashes":["FB7uA1nB4/6wAvcvCTBJ8UcUXgpxoOqe65UgHsU8HUA=", "Oh9qbQ7BtOATK9vy8lXAQZu508niFa+W+tyyl7R+CII=", "g+L+GxNHar/fvsLRu1AiUcQxSyv+9uDscSKYDK2XT6Y=", "6qgYiPJOjr3gN5EaX+3qZ3MuUWE4ItJHTPYICmXELvM=", "pA+KmHqxg0E8DZbdo6oaCl8+lBOOxxcRT/I5X5dXLa4=", "E9khCu+JBWvxI1LYoo5Zu3/u6wB9Ul+z3uVJrDYdnZ4=", "EZ9926bjgfuI5NuGkcblSMVYpjdcbNE9CsHIXV760go=", "smm/CTKGpFV9Ym5y8gc3ioBf689FSSAfENQZaOwK2a4=", "pHtX82Y74bvwhuKZfb907BY0kRutGODi0V+g0AevY3U=", "glyQw5Cb7UTQcVC8Nv6dsXFY/Q41T/OzJMUPByTDcuQ=", "lXysSfZzBRrHF3PliB4b5n2skqb1A6lQIEKJfxyqFu8=", "KToNGOU/2LQPMnwkW0YBQmo27fLOVdrtlPNJA6eS4kg=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201871354\n54qruKG0yzo+f75Kpm/kL2m289lYdeozLCDD9m9SNOQ=\n\n— rekor.sigstore.dev wNI9ajBEAiBY3LhnRJV2GIECj7oMIyNTFVVAuxfsLbMyAf0QLcRzAwIgC86Kt3RkhW8yfsdrfGbRefpKiF43Lb9OiWVcZdtkqh0=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI1ODI1Yjk3MDlkZGZhNWUxYjI3Y2I5NjZlNTAwOTZiODUwNDdlODc0MjgwZjVkYTYyODk2ZDNiYjM3YzMzNmI2In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRXBwV0dNbGgrM1ZKZ2lCWWdEcy8yTXd3ZmJ2dTdKR09PaldidjFIQmlSNEFpRUExRENYNDExdVdHMjFJRFlKQkZ5Q05qeWxKYTFGY1ZrU203VWZURWlVanpvPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KTFowRjNTVUpCWjBsVlVYTlNTVVpwTkhOc2JVdG1lVWhWUTFoalVETlhOSGhJUTBkRmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMVJGZVZkb1kwNU5hbGw0VFVSRk5VMVVZM2RQVkVWNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZzU1U5U05YbHdUVXBYTmtsc1ltZ3ZXak5EVjJaTFVIZHlZa2xvUmpkbldVUTRNWG9LV2tkWFpUZDBWMkpTY2tKSlFsbzNSVU12VVU1R1RYUkVNMjVZZFZWWVdsWmFUMmxXZW0wM09HeGxUMlZTZFhSTGJqWlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKd1lVZFJRVUZDUVUxQlVucENSa0ZwUlVGNE15OWlWa1ZVVkRZNFZsaFdZVEJuZFVkMWJUSmhTR2hEY0Vnd1pEUTRXQXA1U0daSFVtRjZWMFJzYzBOSlFYWllZa1pTYTJodlJrMUlTR2x1TlRad05FUnJMMHQzWW1sSk0wOXFlamh2WjJack5qUkNiM1JZUjAxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKclFVMUhXVU5OVVVNMVpFMU1Ua3MzZDNCVFYwbE1iVk4wUTJWWWJEaGtTVzk0TkZWblIyOTZNa2RTV0ZwbWFURnROMVpzUjI0S01rUmFZV0UzYlZsM1lXVmhkRkU0T1VseVNVTk5VVU42WkVKVU9VMXFlV1V5UmxCMFZIQldjM2d3TXpadmFWRjRZMXBEVFROeFNHWnFlV3d3YVVKVGJ3b3JiVGgzVDI1TmVVZE1lbVZJVTB0MGQxUnZlVUk0Y3owS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgpY03shN3i0hD2L5Mia+LbRZbUg/FAIlwdFBJvIZtqxcCFQDhb/5g26zFl/kPCI1juB4mz2ufeRgPMjAyNjEwMTkxNjU5MTJaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE2NTkxMlowLwYJKoZIhvcNAQkEMSIEIFI6xh9GiPa8s/nXTtp3ouQ6VTu3Iomk154rh5Vy5IfxMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMHGWJz0ewJb82rEVInQh1ntGH6Fa53Sxs+d9Usr2gLwIPxvnN/86eGZiwGWy1HiUQQIwTWUTur9/SfI8jEjir6Oh5s8yMFo3Lkf4G4JA8fNKsKs5TR3Txm+t6GPBWgmV+K8r"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"WCW5cJ3fpeGyfLlm5QCWuFBH6HQoD12mKJbTuzfDNrY="}, "signature":"MEUCIEppWGMlh+3VJgiBYgDs/2Mwwfbvu7JGOOjWbv1HBiR4AiEA1DCX411uWG21IDYJBFyCNjylJa1FcVkSm7UfTEiUjzo="}} \ No newline at end of file diff --git a/build/torch211-cu128-x86_64-linux/symbols.json b/build/torch211-cu128-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch211-cu128-x86_64-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-x86_64-linux/__init__.py b/build/torch211-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch211-cu130-x86_64-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-x86_64-linux/_ops.py b/build/torch211-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch211-cu130-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch211-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..9766ea6ec1462866a6c41a1bc002554395f56034 --- /dev/null +++ b/build/torch211-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:07d3eec7fcc9772fee93a595d55d5728e498491cf8ff9a96b19c8c5ed2cf4794 +size 410664 diff --git a/build/torch211-cu130-x86_64-linux/metadata.json b/build/torch211-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..48be0b513a86b5356509165c277e603c96978e91 --- /dev/null +++ b/build/torch211-cu130-x86_64-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": "B9Pux/zJdy/uk6WV1V1XKOSYSRz4/5qWsZyMXtLPR5Q=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch211-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..ef7dfe395aaceec86dfbb4c1e8eb3826ec6b10a9 --- /dev/null +++ b/build/torch211-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrOgAwIBAgIUOC3/PVMMBgYq+hbAPyg3ey5xxIgwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1OTE1WhcNMjYxMDE5MTcwOTE1WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEvnvv7M1AnoBaEvdy3iIluWx3nYqjO7r6VbQQGLCgTBSBa7b5CBOKGqREyZJAnx6mTganqlVMlH46RUsKYGnyE6OCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRplOQAABAMASDBGAiEA4qBhjGRUxb1pPtsthMF14Whh/yjgfKnBDyLG0waeBnwCIQDOS5vw0GB5Z9S6nMt6oYAtddwtBT0zpfxCgeSIuDd9TjAKBggqhkjOPQQDAwNnADBkAjARNPeM08QOFdPznmi6lBzNG6CmfTZN0ufc4cW+Kzz7XZbGoSV33vjfRMEGR2NQ09QCMHyUYOadA3lR2bBLk1Vuq1HmBmQ76+KKrNNd3qiDUvMgEh0JvlrDcwI1VgQOKYvp7g=="}, "tlogEntries":[{"logIndex":"3323775966", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429156", "inclusionPromise":{"signedEntryTimestamp":"MEQCIBpDZIo04UNs4e1k9R+RRXopWNSCAWAtMupNIvPpDH+DAiAEgxu3t1HARgGnq6BPOharDyagcOOnjO8IzFqTe+9WDQ=="}, "inclusionProof":{"logIndex":"3201871704", "rootHash":"Yz4tugnWyAY7tbWtrwpTKRSJ4tGX5u839YJmDogoVdo=", "treeSize":"3201872002", "hashes":["8YkGQbDNB/14Wst6wZjXX0gPgSvgNJiJQ7OUBwruasw=", "xQj+9nzKciV6/ojpbp2RfC8dzy04EN0zibXxvESA8yI=", "rOOthtnau7N3Ev21vOkND5WiLoSbpvEUc4wDtj+DPgU=", "S3+UxMaK1PKWB/GI+BESCQ87lDRKLRVwJwEysERC8HY=", "joq/l9uUx2wxx61x8eKlE8IDh0Qh0D2L8wMAb11kBac=", "rQJTIqNJl6fkx9Ma9NDNkYFBoWUGsQ47FIr1Y8P41FA=", "qhqlwv/DyfKoRnU/L0bqmJiIpFsHDFxMNIDIDiwrQnc=", "13FvdE1/nw/9bPyQQd98Tssyj8mWxkxECMH91ZIwfqY=", "78trHy97IRNOGtmVsqX/84epZR7RBz0dBWYsoSJxJL0=", "szrJoWEZBlN6vVyqAH7cigkWXjzTEDo9LT8MwP+/1iA=", "glyQw5Cb7UTQcVC8Nv6dsXFY/Q41T/OzJMUPByTDcuQ=", "lXysSfZzBRrHF3PliB4b5n2skqb1A6lQIEKJfxyqFu8=", "PJcEWheSvn2yVXD8wR9GJ6lKvUSCDM+m1WGh/i9UHEc=", "KToNGOU/2LQPMnwkW0YBQmo27fLOVdrtlPNJA6eS4kg=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201872002\nYz4tugnWyAY7tbWtrwpTKRSJ4tGX5u839YJmDogoVdo=\n\n— rekor.sigstore.dev wNI9ajBFAiEAzpa2gTTnt0z/VLUcVwJ6mwRbNXY5tWeEyZk9d7toLPUCIDsZZJ3UWrkIgwBf5YgB2WbQvR292FpdFlR/PbMS+y60\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJkYjY5YWQyNjkyOGM3ZWUzZTdhOWI5ZTFjMTg5NDNmNGRlMzEzYjg2YmJjMWE0NTM2OTMxZWQ3OGNkZmEzYWRmIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNwZDNkVm0zZUJicVUxeXErRzRIZmM5b2c1RVMzY2FDSE5rYlE3T1FSYmxBSWdhWHNiZys2QnJsNmJUVEk2RDRZTlJZb1RIRU40dzBUVnQzWUZ4RVNQVm5nPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KUFowRjNTVUpCWjBsVlQwTXpMMUJXVFUxQ1oxbHhLMmhpUVZCNVp6TmxlVFY0ZUVsbmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMVJGTVZkb1kwNU5hbGw0VFVSRk5VMVVZM2RQVkVVeFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVYyYm5aMk4wMHhRVzV2UW1GRmRtUjVNMmxKYkhWWGVETnVXWEZxVHpkeU5sWmlVVkVLUjB4RFoxUkNVMEpoTjJJMVEwSlBTMGR4VWtWNVdrcEJibmcyYlZSbllXNXhiRlpOYkVnME5sSlZjMHRaUjI1NVJUWlBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKd2JFOVJRVUZDUVUxQlUwUkNSMEZwUlVFMGNVSm9ha2RTVlhoaU1YQlFkSE4wYUUxR01UUlhhR2d2ZVdwblprdHVRZ3BFZVV4SE1IZGhaVUp1ZDBOSlVVUlBVelYyZHpCSFFqVmFPVk0yYmsxME5tOVpRWFJrWkhkMFFsUXdlbkJtZUVOblpWTkpkVVJrT1ZScVFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2JrRkVRbXRCYWtGU1RsQmxUVEE0VVU5R1pGQjZibTFwTm14Q2VrNUhOa050WmxSYVRqQjFabU0wWTFjclMzcDZOMWhhWWtjS2IxTldNek4yYW1aU1RVVkhVakpPVVRBNVVVTk5TSGxWV1U5aFpFRXpiRkl5WWtKTWF6RldkWEV4U0cxQ2JWRTNOaXRMUzNKT1RtUXpjV2xFVlhaTlp3cEZhREJLZG14eVJHTjNTVEZXWjFGUFMxbDJjRGRuUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgNX4BitZqjhR1jzopyD3tA7+UJ2DnrpRMC3/1jI5LSG4CFQCsvm3vjB7mbxjM74rsdUcoxB2x4RgPMjAyNjEwMTkxNjU5MTVaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE2NTkxNVowLwYJKoZIhvcNAQkEMSIEIKF6+bDHzaGPPctlmBCNAudM2/CPdLVXKDdrtofojWA3MIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMGhDPpOs/7Sbf94iNzMKvKFNv/Zlic/vLDWTmn53yASx3wuzoKnxOurGpDXZIxYj1wIwVR3z1lfGRr+7qLqS4fnGKnzY5f8kYUxwoG4D0OgPCNzYOc5F6kPD59RpmPasq5tc"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"22mtJpKMfuPnqbnhwYlD9N4xO4a7waRTaTHteM36Ot8="}, "signature":"MEUCIQCpd3dVm3eBbqU1yq+G4Hfc9og5ES3caCHNkbQ7OQRblAIgaXsbg+6Brl6bTTI6D4YNRYoTHEN4w0TVt3YFxESPVng="}} \ No newline at end of file diff --git a/build/torch211-cu130-x86_64-linux/symbols.json b/build/torch211-cu130-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch211-cu130-x86_64-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-x86_64-linux/__init__.py b/build/torch212-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch212-cu126-x86_64-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-x86_64-linux/_ops.py b/build/torch212-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch212-cu126-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch212-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..1f66c60e52619223d48727bfb14790c081d650f3 --- /dev/null +++ b/build/torch212-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f8c7cfd4d57158ac319babf49b0a661340fc71546abe21d04b7b4a0fc2952294 +size 229392 diff --git a/build/torch212-cu126-x86_64-linux/metadata.json b/build/torch212-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..14f738047c379894fb8dfe87daa0b5cebbe9dfb7 --- /dev/null +++ b/build/torch212-cu126-x86_64-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": "+MfP1NVxWKwxm6v0mwpmE0D8cVRqviHQS3tKD8KVIpQ=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch212-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..1b582f9d10e990e412be2affbed16f2d27609921 --- /dev/null +++ b/build/torch212-cu126-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrKgAwIBAgIUF5MbeIqlrG2NKn/WU4wPv65jujQwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1OTIwWhcNMjYxMDE5MTcwOTIwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEhhljUz2cr+tacQpr4Xut42anto+VHny3mMS7pUCw8jlMP/jHxuO5T6Fhs7bNYRprPTD/OrLN9hVks9cBXwlSOqOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRp3YwAABAMARzBFAiAAp6XtZzdj4PeFa6lGd7N7UtRu27lUW1quyTtC0BN8EgIhANa9TJv1N0bzjeUOZtbCfj7PPaVV4ZjOtiorDuEObeFfMAoGCCqGSM49BAMDA2kAMGYCMQCY7upbhax8Zsv1/qrlnU+cEvQ49MNAPzAlp7sI9t0qyZa6URdLPyofqEi0SN7vPvsCMQDOnezm2NisIAUbMwyMhAxXmqzdZB3L2b/WBrxX2jxFvqmedfKgeMoW2CfNx9gyWd8="}, "tlogEntries":[{"logIndex":"3323776834", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429160", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQCihW0IDwKODLuk6gLbzMZtYM4R72AfI0M7VwSWq+HvQgIhAIHnfecZ1yc+Tekr97tyr7YJ8LG7X/6w/Xi2jIZri6oP"}, "inclusionProof":{"logIndex":"3201872572", "rootHash":"HkfrimUwfeDnAs/vdijQ5P7/5fPAUSo2SvYQvycrZ+c=", "treeSize":"3201873289", "hashes":["D4R5d5Rz3tSh1DeCniYjIMBcfFbC8xh5hYmiQTR+P9g=", "VucSTmOvfwD7+wfqlsQTpglBO8f46tGo+rrd4RIJlB8=", "k2c6OPGAHwgRlBd1fljj1w+ogH52cJPHsWsTEOSQRzg=", "eOPzX5Ap5/GOT7Pq0ecyAfeIcQe2jcKjcDL1Swjv/90=", "ExNLxPvSGfJEaTv6CT/t/cerBOOt49ddrVae6XM8AtY=", "pAz/MnHMrO9ImU424ihIvFmDEtwjbCLrjec5qh3+lSg=", "qF9ZCmV3EFXTFiXuY01grqOeLOoVJOSxO2wPLsRgfD4=", "Z5WfiivB45HwESmKdt0XsM2v2TEOyViqW36qXBSHl+Q=", "GkxOv2b8NufnFgvo65foAOa06g2YfG7XLBRkPhkWkq4=", "JGM1cS7yJFLgFSXinmOuunK+Z07ZVX9Wji+ornJ68Qo=", "TvNaI/CuFix8Tn8ppVllS45JH4WQwPJefwsJXPYHZLc=", "vjPNDK+mGQQHdOhpF1epI0ArReuRdQTWwCGLKzODMXI=", "KToNGOU/2LQPMnwkW0YBQmo27fLOVdrtlPNJA6eS4kg=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201873289\nHkfrimUwfeDnAs/vdijQ5P7/5fPAUSo2SvYQvycrZ+c=\n\n— rekor.sigstore.dev wNI9ajBFAiEA0qcR9I4VR1LnVqaGfhka6MbbU0eA+dfJb4Ctn1+uHqoCICFfcQ0zTNxfIuXFOXB5u3GriVL+U5UpECRKQuk4Sd1B\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIzYzAxYTQ2NTEyYjRhNzY5ZjcxMTAxYjcwYzcwM2FmYWU2NmQyOTY5NmU5YjIxZTc1NWZjOWMzZDFlNzJjNTQ5In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUURYT2dOL29xVElIUXRFRXR6ajZTRW4wWk5oZmk3aDVDZzBxL1hyUHk1ZVN3SWdQeitvZDIyYWtQSjJiS2dweVFhMVlRVCtRRU9NV2Y3ckc4Z3FsdHdHQ1BrPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KTFowRjNTVUpCWjBsVlJqVk5ZbVZKY1d4eVJ6Sk9TMjR2VjFVMGQxQjJOalZxZFdwUmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMVJKZDFkb1kwNU5hbGw0VFVSRk5VMVVZM2RQVkVsM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZvYUd4cVZYb3lZM0lyZEdGalVYQnlORmgxZERReVlXNTBieXRXU0c1NU0yMU5VemNLY0ZWRGR6aHFiRTFRTDJwSWVIVlBOVlEyUm1oek4ySk9XVkp3Y2xCVVJDOVBja3hPT1doV2EzTTVZMEpZZDJ4VFQzRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKd00xbDNRVUZDUVUxQlVucENSa0ZwUVVGd05saDBXbnBrYWpSUVpVWmhObXhIWkRkT04xVjBVblV5TjJ4VlZ6RnhkUXA1VkhSRE1FSk9PRVZuU1doQlRtRTVWRXAyTVU0d1lucHFaVlZQV25SaVEyWnFOMUJRWVZaV05GcHFUM1JwYjNKRWRVVlBZbVZHWmsxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKclFVMUhXVU5OVVVOWk4zVndZbWhoZURoYWMzWXhMM0Z5Ykc1VksyTkZkbEUwT1UxT1FWQjZRV3h3TjNOSk9YUXdjWGxhWVRZS1ZWSmtURkI1YjJaeFJXa3dVMDQzZGxCMmMwTk5VVVJQYm1WNmJUSk9hWE5KUVZWaVRYZDVUV2hCZUZodGNYcGtXa0l6VERKaUwxZENjbmhZTW1wNFJncDJjVzFsWkdaTFoyVk5iMWN5UTJaT2VEbG5lVmRrT0QwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg+NU2n2Rk3d997zkiRvqWwKjnSwJGTzPi7Jiy8OYTnzUCFA+szdclutX2BEgZBYkwgrMfh12oGA8yMDI2MTAxOTE2NTkyMFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTY1OTIwWjAvBgkqhkiG9w0BCQQxIgQgb/mXxw1qMETvfoo+nj49Pzn62p9JqmV20Af1soy3RlkwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIwPYuqFECQ55kaO7wkNlFztH9g7aWiQebhI3q7YqZ69P7SU/p2YrKeBqXBSMtLONrVAjEAnJGJCrxGn+2UU+Bun1wglzSNGCrI/PBqQ0uFFFyU2wyypRQwDU+VBwinZOpgeBuv"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"PAGkZRK0p2n3EQG3DHA6+uZtKWlumyHnVfycPR5yxUk="}, "signature":"MEUCIQDXOgN/oqTIHQtEEtzj6SEn0ZNhfi7h5Cg0q/XrPy5eSwIgPz+od22akPJ2bKgpyQa1YQT+QEOMWf7rG8gqltwGCPk="}} \ No newline at end of file diff --git a/build/torch212-cu126-x86_64-linux/symbols.json b/build/torch212-cu126-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch212-cu126-x86_64-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-x86_64-linux/__init__.py b/build/torch212-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch212-cu130-x86_64-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-x86_64-linux/_ops.py b/build/torch212-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch212-cu130-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch212-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..838791d00a2b02c6cb9eddca69803e01898f1ce6 --- /dev/null +++ b/build/torch212-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3046373069d07d08d045716387c349019d7b9b244e08f860f0a85db04216d34c +size 410680 diff --git a/build/torch212-cu130-x86_64-linux/metadata.json b/build/torch212-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..3f392dfb9889b84a118acdf8381b3ae4696d8802 --- /dev/null +++ b/build/torch212-cu130-x86_64-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": "MEY3MGnQfQjQRXFjh8NJAZ17myROCPhg8KhdsEIW00w=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch212-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..9bdfb99deec498a4b6315e6897dd922f648495b1 --- /dev/null +++ b/build/torch212-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrKgAwIBAgIUQJrc6q1AKXTB4wc1obdd99yeVdswCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1OTI0WhcNMjYxMDE5MTcwOTI0WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE9dN1L/DVmWORZ+jEpQTftzTGumPqJvLTZqnFjfTcfwTssoZNSeKPwBzybm216L+kc3N90AnjFe0tE/NsQh/VbqOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRqJmAAABAMARzBFAiEAyeQqP6Oy4CrE+Lfm43wsx0Yczv2zg6w2W3VB5t8HZ8sCIHDd8zvA6HN65w1iBDFjDwPwW0MLtCt/249/HZ62jdbgMAoGCCqGSM49BAMDA2kAMGYCMQDXckxMzIWiYtuYhZapUCAWHQA/8VcXwPsoIfToQ9G1H9aa+Lvhcwkvo2XKR6Tfx/MCMQDQTfNH5GcQp45yNuE6Zl8jY0tU1kxNJTcx4sfXeUw0FmTt/MSWL4GyBFOUE+6UuVw="}, "tlogEntries":[{"logIndex":"3323777767", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429165", "inclusionPromise":{"signedEntryTimestamp":"MEQCICLSiysLWohBmMHz46EL/XtdNbOJtyFPiOZMZI3MFFpNAiAp4yZuTqncJRiyDaIqCLjcxMd9TK5C/JDzyJlugZbpdA=="}, "inclusionProof":{"logIndex":"3201873505", "rootHash":"/kPB9vAB+R9486Vk4/D+ia3M5kOeVOVs+ExT+ZrUOug=", "treeSize":"3201873780", "hashes":["m1ke0in3wihEH7e52/8GqdQf7F9NG5b9Kcp4OWgLahU=", "MWSTga/NdBbtu/uKy3lSkFXKyVayGkwyEA/eMwWA5Ik=", "1hPX9MvN2DLVZQZvSI5RwGQSE+MuT8RqYgoXHM5BTLI=", "LRJdYkE1zd9lPBZVdhoWDvQ0vHnQMt6ZDHQpB6kus3s=", "Xoa1mgrVAeNKtgTkTJiW+iPaXb7H2tLZHs2BqDiGy1I=", "M2HgEm1jiZNyEnCZOnYQV31Pqn8Ot+fyzWpGTIMRFA8=", "4EXAfyloCuOL1fyRqLNyNuzGSKPf2PSq4wmtIp5nHLE=", "fRDCYkoDkM7PkY7KK2aGCIThuQMJzeGghFi0vwwrPW4=", "IRiTOeBQ0ILSFoZ8k6KgHJ6QY3GCx93oMVUfr85UcwQ=", "p/uvC4RXKhROjBISEKFGTf+jBSWcQJ8d3jmpnP9oMGQ=", "CdOR2r7420N3zLSdpo0AJtgyEHVYzWtl0E7Bws5Zg7w=", "vjPNDK+mGQQHdOhpF1epI0ArReuRdQTWwCGLKzODMXI=", "KToNGOU/2LQPMnwkW0YBQmo27fLOVdrtlPNJA6eS4kg=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201873780\n/kPB9vAB+R9486Vk4/D+ia3M5kOeVOVs+ExT+ZrUOug=\n\n— rekor.sigstore.dev wNI9ajBFAiAKh+8tGhMcfi85cBvVAdhACtWXE0NfGOy+UNDq3xtx1gIhANFlXH3ggIilOirJCc8IpZic37GvhY630Ttxq1HggADY\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJkZDcxZmFmNjkzMTJhNzI0ZDdkY2NkZjc0NzRkMTc0ODNlNDM4ODI3ZDU1OTdkNWUwNmI5ZGFjYzdhYzdkZGVlIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRHVxZ0NzQkM3MkE1bC93Z3krVC8xcm53M1BOWHZLZUtjRmh5eTJ1aWhsYkFpRUE4dXJKRDZIWk1GWHJiZFJNeE1iOFRFNWRETjg4aWpLWlU0NDQxY1ZHM3dNPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KTFowRjNTVUpCWjBsVlVVcHlZelp4TVVGTFdGUkNOSGRqTVc5aVpHUTVPWGxsVm1SemQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMVJKTUZkb1kwNU5hbGw0VFVSRk5VMVVZM2RQVkVrd1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVU1WkU0eFRDOUVWbTFYVDFKYUsycEZjRkZVWm5SNlZFZDFiVkJ4U25aTVZGcHhia1lLYW1aVVkyWjNWSE56YjFwT1UyVkxVSGRDZW5saWJUSXhOa3dyYTJNelRqa3dRVzVxUm1Vd2RFVXZUbk5SYUM5V1luRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKeFNtMUJRVUZDUVUxQlVucENSa0ZwUlVGNVpWRnhVRFpQZVRSRGNrVXJUR1p0TkROM2MzZ3dXV042ZGpKNlp6WjNNZ3BYTTFaQ05YUTRTRm80YzBOSlNFUmtPSHAyUVRaSVRqWTFkekZwUWtSR2FrUjNVSGRYTUUxTWRFTjBMekkwT1M5SVdqWXlhbVJpWjAxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKclFVMUhXVU5OVVVSWVkydDRUWHBKVjJsWmRIVlphRnBoY0ZWRFFWZElVVUV2T0ZaaldIZFFjMjlKWmxSdlVUbEhNVWc1WVdFS0sweDJhR04zYTNadk1saExValpVWm5ndlRVTk5VVVJSVkdaT1NEVkhZMUZ3TkRWNVRuVkZObHBzT0dwWk1IUlZNV3Q0VGtwVVkzZzBjMlpZWlZWM01BcEdiVlIwTDAxVFYwdzBSM2xDUms5VlJTczJWWFZXZHowS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgiHPth2Q1QDqtn6yP89b6eZXoTjxVjuZp8S1bVbgTRC8CFFh+kZLmIutf7+1ylfalSgRcS501GA8yMDI2MTAxOTE2NTkyNVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTY1OTI1WjAvBgkqhkiG9w0BCQQxIgQgwDED44mteQPAPn66IHnGFAJMupbwWgdrslDEI2PWC1owgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAKb7va7CDoiFmEwiQnEpTqyQeGElNQVOyookPfGF6xb6M0SgPjr4YAysZUywAXlVLwIxANTDHwt/bW2T+m3xkKaxKdGMYfBA1R1zCWrnj0KA2YKHh87f45O1FHIW8QTcMdXWow=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"3XH69pMSpyTX3M33R00XSD5DiCfVWX1eBrnazHrH3e4="}, "signature":"MEUCIDuqgCsBC72A5l/wgy+T/1rnw3PNXvKeKcFhyy2uihlbAiEA8urJD6HZMFXrbdRMxMb8TE5dDN88ijKZU4441cVG3wM="}} \ No newline at end of file diff --git a/build/torch212-cu130-x86_64-linux/symbols.json b/build/torch212-cu130-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch212-cu130-x86_64-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-x86_64-linux/__init__.py b/build/torch212-cu132-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch212-cu132-x86_64-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-x86_64-linux/_ops.py b/build/torch212-cu132-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch212-cu132-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch212-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..641c55b5be60bf96b2318deedebda58bd5ccc4fc --- /dev/null +++ b/build/torch212-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6c23d953253e303e9eb13ea7cdee2f3340e47f217c567d555b274c898c10d72e +size 449528 diff --git a/build/torch212-cu132-x86_64-linux/metadata.json b/build/torch212-cu132-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..e3ec087847545b6beb889c06a17c41b3bea68ae2 --- /dev/null +++ b/build/torch212-cu132-x86_64-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": "bCPZUyU+MD6esT6nze4vM0DkfyF8Vn1VWydMiYwQ1y4=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch212-cu132-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..aa9eeddb9ae2565a9a7b96a1b219884458b4c14c --- /dev/null +++ b/build/torch212-cu132-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHKzCCBrGgAwIBAgIUXborIMcSzvDqHdKzp8a64i2gBbYwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1OTMyWhcNMjYxMDE5MTcwOTMyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE/cbM4UqqFluWAGcQG1F9F48Cn0x3li4jwEXFUBM54MMiU0W0HOrbXm81A8TsiFTUBsssKN5xIHnvAgEAymrnTqOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRqlmQAABAMARjBEAiBd1X+vVl1ZH+f3cgoAJEnBMeXqZrKCpjOrJEJuxpMSsAIgcSoEVPbiLA2h42j68WPjusVJdsgFV0YVq9Lw49Ds19wwCgYIKoZIzj0EAwMDaAAwZQIwcxEBrCMy9bVkMNFJRpHDSwSjmD6Anay/ICDDDZqzY5vUs+CW133P5UO1sEWxbNfOAjEA0Wkgek85quwPZJjpzROqNSwarF2wTBlFnADnrqQBCJLFwjhKwNsNBfhDdc3QHSWG"}, "tlogEntries":[{"logIndex":"3323779259", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429173", "inclusionPromise":{"signedEntryTimestamp":"MEQCIE2Hp4b+NZeRhTCo5TXgSqyWfgpeNjWicaQm/gl2aiJeAiA1HO8VrHKUpDClUdPMvUu8xJoVxUq1mhjYTMJpi9K94w=="}, "inclusionProof":{"logIndex":"3201874997", "rootHash":"Ri6UImLnc0IGb55b78VvDj5NgbkO1CunG/whjfwMQKg=", "treeSize":"3201875226", "hashes":["u7OQD3MQOcjPoMJWa/glANe+T6VuEUEiMEThCBD73rc=", "qs43WxJ6uZBD54O74tOCnkR8Gn0pP2EpJ++7XqNMB7o=", "CVL2pATNVK3VQfT21mEH+9AFfjM4YUijB1DzPby5EaY=", "b3Z7Gwty0qULtt1pf1FH9T9SsdJtX3H018SU24XFYa8=", "/6LmAZZGtlSlGW1legSv0Pf3//6VWntY78sic0G2rY0=", "yAROlkGwEv07tEM7H8x58+7T2uX8iw6YejDWpcI0OHU=", "1L78mEgHdnkOI17MFg7ezIWD4BnCLClkEOcRSrbXncE=", "02j/uuCSi1BMTM0cUZRxkLZGY9q4Xra09auw94aWxlc=", "jA9DRlU+OxTeDq5jZkB0O8E7am/OqXCOsCNqxnGlxkU=", "4mn2xec57v/mZms/2v+2I9CmAUojj15NwyZgThFApV8=", "PU2pUKT7IIwavBppqRYwhfVEenuVPz+z0PSa4KHihwg=", "vjPNDK+mGQQHdOhpF1epI0ArReuRdQTWwCGLKzODMXI=", "KToNGOU/2LQPMnwkW0YBQmo27fLOVdrtlPNJA6eS4kg=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201875226\nRi6UImLnc0IGb55b78VvDj5NgbkO1CunG/whjfwMQKg=\n\n— rekor.sigstore.dev wNI9ajBGAiEAmqIjlaxr1x651UzjhBFxu/K0AF6Oc4QWwOXK0Drc8uoCIQD+X6JUY1SoLkhQuJBysdxkxjDVYqcIqZPP0nbQe0iLFw==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIyM2MwNzA3Y2UwYzZmYTE2ZTU4Y2UxM2U0ZGRjNTdjYzcyMTgzYWViY2ZmMzQzM2M0MTBlZmUwMzNlNDVjNjhhIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNBMTRVeXZ1QTM5WGUzamZYVndPS09ZOEJhQ2dHYnFxZmNJd01udTJodUd3SWdiL1A1MjNJMXFhWEhhdTR2QnpoaWlGUUdQMXM5Mnk2bWNNN3pVbTVVTjAwPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExla05EUW5KSFowRjNTVUpCWjBsVldHSnZja2xOWTFONmRrUnhTR1JMZW5BNFlUWTBhVEpuUW1KWmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMVJOZVZkb1kwNU5hbGw0VFVSRk5VMVVZM2RQVkUxNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVV2WTJKTk5GVnhjVVpzZFZkQlIyTlJSekZHT1VZME9FTnVNSGd6YkdrMGFuZEZXRVlLVlVKTk5UUk5UV2xWTUZjd1NFOXlZbGh0T0RGQk9GUnphVVpVVlVKemMzTkxUalY0U1VodWRrRm5SVUY1YlhKdVZIRlBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKeGJHMVJRVUZDUVUxQlVtcENSVUZwUW1ReFdDdDJWbXd4V2tnclpqTmpaMjlCU2tWdVFrMWxXSEZhY2t0RGNHcFBjZ3BLUlVwMWVIQk5VM05CU1dkalUyOUZWbEJpYVV4Qk1tZzBNbW8yT0ZkUWFuVnpWa3BrYzJkR1ZqQlpWbkU1VEhjME9VUnpNVGwzZDBObldVbExiMXBKQ25wcU1FVkJkMDFFWVVGQmQxcFJTWGRqZUVWQ2NrTk5lVGxpVm10TlRrWktVbkJJUkZOM1UycHRSRFpCYm1GNUwwbERSRVJFV25GNldUVjJWWE1yUTFjS01UTXpVRFZWVHpGelJWZDRZazVtVDBGcVJVRXdWMnRuWldzNE5YRjFkMUJhU21wd2VsSlBjVTVUZDJGeVJqSjNWRUpzUm01QlJHNXljVkZDUTBwTVJncDNhbWhMZDA1elRrSm1hRVJrWXpOUlNGTlhSd290TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgdWYu6VA0wu4b9V+dVTjfPZ/mJtyKgcwkLyVXg0y0L3sCFEniQorbGvsvgFUTivc4eHh/e29uGA8yMDI2MTAxOTE2NTkzMlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTY1OTMyWjAvBgkqhkiG9w0BCQQxIgQgevCdLAuiAg5mUTdLQpt59DCO8UWFTbq28UNMbaKiMpIwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIwOPltB/kb7COndRnxB2Vfn7+REnQRBEFztLZdB9R5lt8OslSRKWaPH5j1HIW4X5uLAjEA4qpHxdZ+y2pLeIEDSKInT7q4Pv+abpdMo3uMgZBpWDzyA4oDFari4Q4e8WhmoiIG"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"I8BwfODG+hbljOE+TdxXzHIYOuvP80M8QQ7+Az5Fxoo="}, "signature":"MEUCIQCA14UyvuA39Xe3jfXVwOKOY8BaCgGbqqfcIwMnu2huGwIgb/P523I1qaXHau4vBzhiiFQGP1s92y6mcM7zUm5UN00="}} \ No newline at end of file diff --git a/build/torch212-cu132-x86_64-linux/symbols.json b/build/torch212-cu132-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch212-cu132-x86_64-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-x86_64-linux/__init__.py b/build/torch213-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch213-cu126-x86_64-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-x86_64-linux/_ops.py b/build/torch213-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch213-cu126-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch213-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..0e1abd46a5217db4657a0edb7c9272560e254362 --- /dev/null +++ b/build/torch213-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3af7023240fc3ca6af9462810cad77d56ae1c34d0303d593c59070bd09a3ac41 +size 229392 diff --git a/build/torch213-cu126-x86_64-linux/metadata.json b/build/torch213-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..0fb237a745eeffeeb0233bf803072d0c3f1987f5 --- /dev/null +++ b/build/torch213-cu126-x86_64-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": "OvcCMkD8PKavlGKBDK131Wrhw00DA9WTxZBwvQmjrEE=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch213-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..e98f8e2cc31674d72dd19f7c0aa4c454b6f4e653 --- /dev/null +++ b/build/torch213-cu126-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLjCCBrOgAwIBAgIUPsh4nGYvAZangJPKuxX/R81FF10wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1OTM2WhcNMjYxMDE5MTcwOTM2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEnzvez96tGqSwUMpK78/NfLO5v9qTbRoNsO0pF6c1XYrtZ1h6kvvGMTthQxG+0u4hBZHQGybMw52n74hLnJw36KOCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRq10QAABAMASDBGAiEA8kZfNzYlzQW0Q2OTXgNHzkTa/tvDS4ohdCnF6ZjL2HgCIQDUyLR7eo/lj39J6cd91MHl41DNAYgrrFGIjySuNRDcoDAKBggqhkjOPQQDAwNpADBmAjEAm7Guu1cr0SIgcnIkXBB6Xdq0xhfF8EAfrYOGy0/7RCC93hahuLMYRhD3QevSLHj5AjEAi8zp4jj/8c4LVW16KrvEun3iMzl4li2ML/H13y+9pLXtKJggubfj9skcrlkzyniM"}, "tlogEntries":[{"logIndex":"3323780001", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429176", "inclusionPromise":{"signedEntryTimestamp":"MEUCIDcP5zmi/nZ18chXr7U/eGSM4JTbaEAoxWVJtUs3RS7HAiEA/TfJI7/SYtVV23HCBRV3MLeTwaPsDIHX+YQKnpYS+YA="}, "inclusionProof":{"logIndex":"3201875739", "rootHash":"BAiQqJF26aPdQ8+el9ADWTb8VrWga35L68fTPf3KiI0=", "treeSize":"3201876383", "hashes":["/fLrT2gl2x+8fx8KXlx+H+4eOch9WatijX0yvvh2/aw=", "/Lvzbzpw08aTVtQsqtwRP8XWltkDjEy3JQ2NoeacPEg=", "mB/GhVGmmbTbnw1f+Vx40D6eRvKZkMHgLq5a4pK6/QU=", "6twgxs1O9zMpMFzldKn58QihMh4HnrGgwb8ixUlbvOk=", "YFKuT5eIC4TSryoZxdIxvHfv1mIfMMG5BwJJngzLCI0=", "VBqEE8C4PlYpZT90yceJz016xudrlG64gbeTpqwNgqg=", "UYhDgb2rcO9IKwnQww2fxDwy8+qTHxxIarYcw+yEI1Q=", "pZISg+gfI6OuX1dcNrqBFA2WtPLfpxVY78MskawAxSk=", "ZJhml+vHu8rSNOTrD7RC++nYNIRT3s87AzvdmIyM1LM=", "GLGVhkmsKW8bf+qt0YwTT9QRMhu+ifh5xFJO1A03q2Q=", "4mn2xec57v/mZms/2v+2I9CmAUojj15NwyZgThFApV8=", "PU2pUKT7IIwavBppqRYwhfVEenuVPz+z0PSa4KHihwg=", "vjPNDK+mGQQHdOhpF1epI0ArReuRdQTWwCGLKzODMXI=", "KToNGOU/2LQPMnwkW0YBQmo27fLOVdrtlPNJA6eS4kg=", "jR1JpLECPELqfHlWoTD1g/yJeWoPWTzQSAgBf23nAGY=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201876383\nBAiQqJF26aPdQ8+el9ADWTb8VrWga35L68fTPf3KiI0=\n\n— rekor.sigstore.dev wNI9ajBFAiEAlA7bEIff8E1igzvAHF5bYsjZTiJ9OENjsmao2qv5COYCIHAQMR6B/LT8TNZg/wc2+VjYnLvollOXaOkMZP0YMABE\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIzMTc3N2I3M2E0OWY4M2EwNjE0ZmE0MGIwMjY4NzZjZjExYjA4OWQ4MmIxMTQ4MzQ2YzQ2NzZkMmZlMGYxYWQ0In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNFZVlFa3VHazdFZjBvM2t1TndiYUI2dnBzZnVQWXd2QmV5dUpKMlFpbEpRSWdQRUxJaUpYN2p1UE5ucGhDQjZaYXlUKy9Zbk4ybVhtOSt6dW50UUJXYTBvPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1ha05EUW5KUFowRjNTVUpCWjBsVlVITm9ORzVIV1haQldtRnVaMHBRUzNWNFdDOVNPREZHUmpFd2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMVJOTWxkb1kwNU5hbGw0VFVSRk5VMVVZM2RQVkUweVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ1ZW5abGVqazJkRWR4VTNkVlRYQkxOemd2VG1aTVR6VjJPWEZVWWxKdlRuTlBNSEFLUmpaak1WaFpjblJhTVdnMmEzWjJSMDFVZEdoUmVFY3JNSFUwYUVKYVNGRkhlV0pOZHpVeWJqYzBhRXh1U25jek5rdFBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKeE1UQlJRVUZDUVUxQlUwUkNSMEZwUlVFNGExcG1UbnBaYkhwUlZ6QlJNazlVV0dkT1NIcHJWR0V2ZEhaRVV6UnZhQXBrUTI1R05scHFUREpJWjBOSlVVUlZlVXhTTjJWdkwyeHFNemxLTm1Oa09URk5TR3cwTVVST1FWbG5jbkpHUjBscWVWTjFUbEpFWTI5RVFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2NFRkVRbTFCYWtWQmJUZEhkWFV4WTNJd1UwbG5ZMjVKYTFoQ1FqWllaSEV3ZUdobVJqaEZRV1p5V1U5SGVUQXZOMUpEUXprS00yaGhhSFZNVFZsU2FFUXpVV1YyVTB4SWFqVkJha1ZCYVRoNmNEUnFhaTg0WXpSTVZsY3hOa3R5ZGtWMWJqTnBUWHBzTkd4cE1rMU1MMGd4TTNrck9RcHdURmgwUzBwblozVmlabW81YzJ0amNteHJlbmx1YVUwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgz29FL7VTtEJ712LTzm1MPb/nshixhq+Ug64P4mK9RxcCFB5klOIKoVpB2MSWI7IXYHCJncmLGA8yMDI2MTAxOTE2NTkzNlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTY1OTM2WjAvBgkqhkiG9w0BCQQxIgQgsNvcOnkofoo3FXs1M0lv21vhijXAW80yU1LC9y0OhgEwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIwFYPB7n03uBRpJIr9sR0+nKmquYu+VYiq2DN5h3Gil2WZOcpgew/IklrjTVhp30XqAjEA2i6fIMhSJ9yKzgCnzZAbbgND2zsEYENh8iwuviPwGvDDOMfciEX0QTsU/sJ+a+OC"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"MXd7c6Sfg6BhT6QLAmh2zxGwidgrEUg0bEZ20v4PGtQ="}, "signature":"MEUCIQCEeYEkuGk7Ef0o3kuNwbaB6vpsfuPYwvBeyuJJ2QilJQIgPELIiJX7juPNnphCB6ZayT+/YnN2mXm9+zuntQBWa0o="}} \ No newline at end of file diff --git a/build/torch213-cu126-x86_64-linux/symbols.json b/build/torch213-cu126-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch213-cu126-x86_64-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-x86_64-linux/__init__.py b/build/torch213-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch213-cu130-x86_64-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-x86_64-linux/_ops.py b/build/torch213-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch213-cu130-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch213-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..824b1426a0a6bd38a84a8e711f90741c06b2aab6 --- /dev/null +++ b/build/torch213-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1316b4f2e8cd5ff4e173e7cddd94f6f8cc114debc6bb3c9f344092a8c593cc5f +size 410680 diff --git a/build/torch213-cu130-x86_64-linux/metadata.json b/build/torch213-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..4cd2a66114b974f241b8499516ec0e662b6a1e82 --- /dev/null +++ b/build/torch213-cu130-x86_64-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": "Exa08ujNX/Thc+fN3ZT2+MwRTevGuzyfNECSqMWTzF8=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch213-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..b623cf4c1569a504daf5ec8a9cdbae742c626716 --- /dev/null +++ b/build/torch213-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrKgAwIBAgIUU1E03yf/lvruZn8frQqHglXtKoowCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1OTQwWhcNMjYxMDE5MTcwOTQwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE0wW7wx368IVJQ2s2neEfY9Zaf5c9c0jKDYQg0rs+s04rGHlfyifEYzoRjbV9v5VJfSlY4zLzRo/yNgUpX7Ff9aOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRrEuQAABAMARzBFAiBgW5nF3MdiQbL6QTV+ECTTNYtT4Mh56QUWpkKuv7IsrwIhAKbMM+6OLvBbDXUuKWejhg8dRHj6P7k7f1YMIBeR/5FjMAoGCCqGSM49BAMDA2gAMGUCMQCuQV/Nqu7KUZEaLlVMyLf1fg5Rjb6agPtfS5UrdOry4hk1ibvrItqhRKg5ymioxxMCMCETXsWWFgoEYfS56MBcPD72IgRr0gmVGjNqMOfwxnU/Hi/048Tju6vE/7gTwk7MEw=="}, "tlogEntries":[{"logIndex":"3323780935", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429181", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQDidyd1Y1K4C5LsGqmLIx50473RKiLFjtaEkwsvdFLJGgIgaWzJONrJ6i7DrIH1Icsj0oeo7wnIrVk5Ug310ee8TCc="}, "inclusionProof":{"logIndex":"3201876673", "rootHash":"QT2274YeOioMEAx/s1032G5M12c/TQE5wzFpGBoSppw=", "treeSize":"3201876755", "hashes":["RNtV0NHWojMHuL85wl2vIg/rwTjhSsokC8a0E2urdrE=", "gTHIIXkNtwJc1MRYZ2Afs37/p55HdZ4Aw0l1dNANk8A=", "MOI3j1bTAE9YzK+ryoMzASxd5EG069F87ZqW+tM3eQQ=", "a9dLgTUNDYguFeHWhUojyS37QkxEGz6xx3mw6yBnwIQ=", "3pS9BaLrZbB1XmLToiRX7u0Oq1krOGdqRThX9x1TafQ=", "oES0EJknIim5HbtpeO5GQdB5xWMvHoyRH3ipoAJ8qtY=", "GC5CugBYdn6eYvtjennjSVtrDAAdHSV/vDmN3+yvwIU=", "yxPyC5ORuYyFlcHhbzkHrkkEjA2aEThbq4RrgnHOeu0=", "D3wImw3yDN7HCMgJLUgVXQoA2VbZioSRwRzQYbQKvuU=", "+EnCdBicMIPdaLS4T93icjmKzIV/2nA10KItY2KfgbY=", "qttRZVUdmA9Dz2AmZd4s0AFXdh4VB0rRnYnOMWDAB9o=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201876755\nQT2274YeOioMEAx/s1032G5M12c/TQE5wzFpGBoSppw=\n\n— rekor.sigstore.dev wNI9ajBEAiAuGpxI/WzNP2wzXtCFmDC7y6reT2iPfkxgyx4p4W3MiAIgfx40XzlStf/PoOeA7+081ku25rfGtxK+mdohIlsiRwo=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjZjU5NTdjYzg5NjNhNDAyOGI2N2EzY2ExNmEyZGZiMGIwMTU1ZDVlMTM1N2U2YjYxMGQyOTg5MWUzMzJkZGQwIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJR2Z6eDlSdldkM01rckhnb2FxVWdmMDNQZ0l0NDlwSlpvVjdKVXdNV1NIdEFpQWZ3ejFoMEMySlJqemtkK3NWRTBDMnFKdUdpN3FPUTYzbGlUREJuWFdjd1E9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KTFowRjNTVUpCWjBsVlZURkZNRE41Wmk5c2RuSjFXbTQ0Wm5KUmNVaG5iRmgwUzI5dmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMVJSZDFkb1kwNU5hbGw0VFVSRk5VMVVZM2RQVkZGM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVV3ZDFjM2QzZ3pOamhKVmtwUk1uTXlibVZGWmxrNVdtRm1OV001WXpCcVMwUlpVV2NLTUhKekszTXdOSEpIU0d4bWVXbG1SVmw2YjFKcVlsWTVkalZXU21aVGJGazBla3g2VW04dmVVNW5WWEJZTjBabU9XRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKeVJYVlJRVUZDUVUxQlVucENSa0ZwUW1kWE5XNUdNMDFrYVZGaVREWlJWRllyUlVOVVZFNVpkRlEwVFdnMU5sRlZWd3B3YTB0MWRqZEpjM0ozU1doQlMySk5UU3MyVDB4MlFtSkVXRlYxUzFkbGFtaG5PR1JTU0dvMlVEZHJOMll4V1UxSlFtVlNMelZHYWsxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKblFVMUhWVU5OVVVOMVVWWXZUbkYxTjB0VldrVmhUR3hXVFhsTVpqRm1aelZTYW1JMllXZFFkR1pUTlZWeVpFOXllVFJvYXpFS2FXSjJja2wwY1doU1MyYzFlVzFwYjNoNFRVTk5RMFZVV0hOWFYwWm5iMFZaWmxNMU5rMUNZMUJFTnpKSloxSnlNR2R0VmtkcVRuRk5UMlozZUc1Vkx3cElhUzh3TkRoVWFuVTJka1V2TjJkVWQyczNUVVYzUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgf/6INvndPaXHLL/3nNhLq71/IS+jfnvOZFG0sLCRDaUCFQCtnI0/q5x7U4r3fC6z/N3azwNC0RgPMjAyNjEwMTkxNjU5NDBaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE2NTk0MFowLwYJKoZIhvcNAQkEMSIEINWk7/lUYWGuKugF48yaq+zsBabaEuAXFF8vkfevfj4MMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQCrzKGSVrNUSHyNZwgDejD6Inm3Z5eg48Dl1GFMSKTNzQJoVpUOGPCEixQ8A/jrBtUCMQC1UiuPWvJwWlxa2DzuiVfaJC1RIPpZyC0/vusLWItUMKS7oThLYFluaVjShnk0fS4="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"z1lXzIljpAKLZ6PKFqLfsLAVXV4TV+a2ENKYkeMy3dA="}, "signature":"MEQCIGfzx9RvWd3MkrHgoaqUgf03PgIt49pJZoV7JUwMWSHtAiAfwz1h0C2JRjzkd+sVE0C2qJuGi7qOQ63liTDBnXWcwQ=="}} \ No newline at end of file diff --git a/build/torch213-cu130-x86_64-linux/symbols.json b/build/torch213-cu130-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch213-cu130-x86_64-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-x86_64-linux/__init__.py b/build/torch213-cu132-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch213-cu132-x86_64-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-x86_64-linux/_ops.py b/build/torch213-cu132-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch213-cu132-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch213-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..68c14c23a18a57f80b1ec36996e181b1cd588553 --- /dev/null +++ b/build/torch213-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fd2db4525dbb5ecbc718c93a0a87cb8c4161300b74fa26d75c58b00411e4f789 +size 449528 diff --git a/build/torch213-cu132-x86_64-linux/metadata.json b/build/torch213-cu132-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..02a9583473920e22023db85ec38502195d866350 --- /dev/null +++ b/build/torch213-cu132-x86_64-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": "/S20Ul27XsvHGMk6CofLjEFhMAt0+ibXXFiwBBHk94k=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch213-cu132-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..c5be148806e3ccf3dcb1f77a48f160e14b21e56c --- /dev/null +++ b/build/torch213-cu132-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHKzCCBrKgAwIBAgIUSvCxB9Ytea8/3i4j895oXnz+1qwwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1OTQyWhcNMjYxMDE5MTcwOTQyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEO9nw+g0CcG4TkiW2FcmzWhzju/I01PRy88+e9XMerr7/noD1j/om1FTNTz4DVVM/TRjKPsGibjRAVx9DRrgqU6OCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRrPCQAABAMARzBFAiEA3dX8zIITAnLwRIYWCNmItUgEbFRN9KjL8tZ6FzjzJIACIFYNWL8QAUe/HCfubmnMhsIkaiPQcqCjHKtfVOi9easMMAoGCCqGSM49BAMDA2cAMGQCMASX2WKc93QZExjElf+DPX99gCN2GSqKoiA7ktbJolCzf1VTyxD+MrgWbmyCYeIaOQIwJ6FEq8bU+AKOtU28ftuP7dVlQxFX2WYKcXFr2wQuIBpxa/oUujNiaMmebGpLdFeX"}, "tlogEntries":[{"logIndex":"3323781339", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429183", "inclusionPromise":{"signedEntryTimestamp":"MEUCIANoi8s0DG7Y6Vfay5VSP8wwVIDiKGRnHkVYXH4sHD9VAiEAwBiFPxZ+qGnfd5gCMOussWwY01xB+TTt6mj/75FgxxQ="}, "inclusionProof":{"logIndex":"3201877077", "rootHash":"D3pltdoh75JeIZdA4+qH7V1pVBdHitxBeK7VqdSTVDY=", "treeSize":"3201877109", "hashes":["yhTkgU2SQv1/WAALXhQWFph8hSNKl5lmumBGGxCHChw=", "6vrn6u2PpfDRUiBpaLDvshB8wOUpVtUuLQ0oobHWbXU=", "2HOfqOH4hE9nxGu538eSTB6TvUyCYOd7VUEf6G45V2g=", "8caXN57VPdiRmadeTPOrGQ1yelJBOA2BTkpG7IKdyGQ=", "8VdXwfIJi/aVgLEDRWUptjlZ0pdCEwT1tyJw9C6jJOU=", "BuPvulIief0bFpUX4yJVb+brYy8BCeobcKncr7dP5DA=", "Qp9NUhyUvaahEJr9NnHFBe2+NWSaLvh5D/DnihWckMg=", "q4GOtr54zwmA/KNOY7BlgwXtLltHjxZodbuhEfVMpqg=", "qttRZVUdmA9Dz2AmZd4s0AFXdh4VB0rRnYnOMWDAB9o=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201877109\nD3pltdoh75JeIZdA4+qH7V1pVBdHitxBeK7VqdSTVDY=\n\n— rekor.sigstore.dev wNI9ajBGAiEA6qbJ/GpwJwAI4gjGxwzn2vMIuEfb8zjbh5OCjbtvh2QCIQCDWINIIsrH1q2AG7xpZeMPINTgnebP5In+92h/b+Y+aA==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI1NWNjYjgyOGQ5MDM3ODRhMjc5ODcyYzRjYmI3N2E1ZDQyZTRhZWQ4YzdlZjgxYTNjNmFjMGQ2ZDE0NzcwNWNmIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJSFlYZ0FLRUlhRHl2andtMjRwbDJSRzBhUUpsQ3EzZG84RXJyUlFMclB2Y0FpRUFwL1dkNnM4ekxxcUliVDZDSmNPRkE2V1BDQnlrblF6QjkrLzg2Njc4RTMwPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExla05EUW5KTFowRjNTVUpCWjBsVlUzWkRlRUk1V1hSbFlUZ3ZNMmswYWpnNU5XOVlibm9yTVhGM2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMVJSZVZkb1kwNU5hbGw0VFVSRk5VMVVZM2RQVkZGNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZQT1c1M0syY3dRMk5ITkZScmFWY3lSbU50ZWxkb2VtcDFMMGt3TVZCU2VUZzRLMlVLT1ZoTlpYSnlOeTl1YjBReGFpOXZiVEZHVkU1VWVqUkVWbFpOTDFSU2FrdFFjMGRwWW1wU1FWWjRPVVJTY21keFZUWlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKeVVFTlJRVUZDUVUxQlVucENSa0ZwUlVFelpGZzRla2xKVkVGdVRIZFNTVmxYUTA1dFNYUlZaMFZpUmxKT09VdHFUQW80ZEZvMlJucHFla3BKUVVOSlJsbE9WMHc0VVVGVlpTOUlRMloxWW0xdVRXaHpTV3RoYVZCUlkzRkRha2hMZEdaV1QyazVaV0Z6VFUxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKalFVMUhVVU5OUVZOWU1sZExZemt6VVZwRmVHcEZiR1lyUkZCWU9UbG5RMDR5UjFOeFMyOXBRVGRyZEdKS2IyeERlbVl4VmxRS2VYaEVLMDF5WjFkaWJYbERXV1ZKWVU5UlNYZEtOa1pGY1RoaVZTdEJTMDkwVlRJNFpuUjFVRGRrVm14UmVFWllNbGRaUzJOWVJuSXlkMUYxU1VKd2VBcGhMMjlWZFdwT2FXRk5iV1ZpUjNCTVpFWmxXQW90TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgEceLibwlIRv1D1dAmmDrmhgMb5ixvbAkIL6QKbGlW0YCFQCnQjcFkIZsrF8FnjQkImnDW9vbLhgPMjAyNjEwMTkxNjU5NDJaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE2NTk0MlowLwYJKoZIhvcNAQkEMSIEIOU6yilb8nzlDrB3gWEv5J0+WBjyCd9rEKVOSNRO/w0YMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMHs2YQslnCWY+zZYg6PJevBwysvK7GnMdozXbJug2gRVLRbzNJNGRxjqAmQk0FOjgAIxALq/t/0y43ZwHf7IXcl/o36/+SitWhKRQB11PkOvOAIA8yVbl8sXNDYyjjg1mPHZWQ=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"Vcy4KNkDeEonmHLEy7d6XULkrtjH74GjxqwNbRR3Bc8="}, "signature":"MEUCIHYXgAKEIaDyvjwm24pl2RG0aQJlCq3do8ErrRQLrPvcAiEAp/Wd6s8zLqqIbT6CJcOFA6WPCByknQzB9+/86678E30="}} \ No newline at end of file diff --git a/build/torch213-cu132-x86_64-linux/symbols.json b/build/torch213-cu132-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch213-cu132-x86_64-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-x86_64-linux/__init__.py b/build/torch214-cu126-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch214-cu126-x86_64-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-x86_64-linux/_ops.py b/build/torch214-cu126-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch214-cu126-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch214-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..1de2b0940ec20ff0809fe708dd18d1bc843e4cc0 --- /dev/null +++ b/build/torch214-cu126-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dfb6ab5c0ab2c7cbef2c265b5fa7b3f1ae69f3d1870823a9c64f41e0497dccb4 +size 229272 diff --git a/build/torch214-cu126-x86_64-linux/metadata.json b/build/torch214-cu126-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..15afa56586cf004a183f63370a2e872a91879950 --- /dev/null +++ b/build/torch214-cu126-x86_64-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": "37arXAqyx8vvLCZbX6ez8a5p89GHCCOpxk9B4El9zLQ=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch214-cu126-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..286a16b28934454b68a787cb037d48a4e3f7dc73 --- /dev/null +++ b/build/torch214-cu126-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrGgAwIBAgIUGKwX0C1koTG+Km+Ch3QS7EhlLK4wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1OTQ3WhcNMjYxMDE5MTcwOTQ3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEXGcAORsKya7Q4rQ1G3mbYYMwA8IyoneX3md57LupfoxxaxGb/s6ZX8OjmFFNKhNEjh1gK9ovY8Xc2+To4LRMrqOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRrhAAAABAMARjBEAiBYGKGNwXyW0ZNQpTAL2HoiG0wiLez1chQfzPg8bhMe1wIgTZC1degzbaN7s12LmylOq34GME2d3Xt56oWLgHlON6UwCgYIKoZIzj0EAwMDaQAwZgIxANUEpHMShKseModnphpcQnSv4MpO0Qzpy+W576rLWl8dPIk0kbmTjprs91obP/roXQIxALlY+c6morKIqhIFTUkxoRLykNXSYVDZj7fVr3G4i015DxoT/B708MbFpG5uQSLmIw=="}, "tlogEntries":[{"logIndex":"3323782301", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429187", "inclusionPromise":{"signedEntryTimestamp":"MEUCIDAzM9aX4qaAFLDC8H1eU6ybqFUiZCVjz6rVMOs/ZMwkAiEAmUFxSCjTsJgHPkeQCd5Bzyc7hEX3v16HLwjZFIzGUEc="}, "inclusionProof":{"logIndex":"3201878039", "rootHash":"6WPkPOQwaRM4tRkxFE+/hLo6j7kQ1vCy9ljllH4vvCQ=", "treeSize":"3201878575", "hashes":["uyNofgkKVlB842gJ/lhQIJK9cpPKbWAsarL2WDRzPFE=", "wbK8z1j/xSbh++Vn1WLP44Uj3DNQ5+GVOBqvuTNoJ+8=", "eCUiVJtxeCMpUQZO3BeZtOxvS4nuc2kDkqFzTINdGPw=", "taAUp/0lV4Pj+BqU1XJeQ1xpDYxfMMEzeCdw9H1flUE=", "7XDjAebyC6F/nwRUkzCxnUXa1pJHcVEMgDiMOvqPE3E=", "w3m8l4YV0iA53K8UpuF7f1PiOdwhQsZMdEVHFdrUGVo=", "hFYWQK8QcOKS3TF8HIyX6NHnNb8zJFYZt0mJvOM+SKg=", "mEKut9p+f3Ehhn0k9BI6mcZbetPD3gRtD59ndeRRfjo=", "tO8a2Nd+Qu412odsUz0yhz8+dOjJepcFPpppVXV5YWI=", "7JQAk+54BBWS7ypTd2wow5NtMSE8nqJvSnvZoi/GJS4=", "PywsDAcGw1w2GA2yax7XMnye0HGAwlZxH8di0wVvUY8=", "qttRZVUdmA9Dz2AmZd4s0AFXdh4VB0rRnYnOMWDAB9o=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201878575\n6WPkPOQwaRM4tRkxFE+/hLo6j7kQ1vCy9ljllH4vvCQ=\n\n— rekor.sigstore.dev wNI9ajBGAiEAx+bnMVClhQK1hRCZnwES+6fkJDSzRrfny+KUBDOFIIgCIQDEWstIWrbLRBlsh2GOoAOvrPKKHxpl6rPxWoBVmcgggQ==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI5MTgzMzliZmQ0YzJlMjYyODQ5NzU1MDc3ZTQyZGU2NGUyMmMwOTllYzEzZjUxMzM5YzAzMDljNmUzYzc5ZWE0In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNMVWNkMm8wWlk2a2ZMYXhxV0tvLzRlVnRoN2hVMUVwR3dQSWhSYXBWMXV3SWdMbmI2ZS9oYm1DeW5CVmh6d0VEMU5hajFkL0k4cmpjS0gwTFRkcWdleElFPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KSFowRjNTVUpCWjBsVlIwdDNXREJETVd0dlZFY3JTMjByUTJnelVWTTNSV2hzVEVzMGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMVJSTTFkb1kwNU5hbGw0VFVSRk5VMVVZM2RQVkZFelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZZUjJOQlQxSnpTM2xoTjFFMGNsRXhSek50WWxsWlRYZEJPRWw1YjI1bFdETnRaRFVLTjB4MWNHWnZlSGhoZUVkaUwzTTJXbGc0VDJwdFJrWk9TMmhPUldwb01XZExPVzkyV1RoWVl6SXJWRzgwVEZKTmNuRlBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKeWFFRkJRVUZDUVUxQlVtcENSVUZwUWxsSFMwZE9kMWg1VnpCYVRsRndWRUZNTWtodmFVY3dkMmxNWlhveFkyaFJaZ3A2VUdjNFltaE5aVEYzU1dkVVdrTXhaR1ZuZW1KaFRqZHpNVEpNYlhsc1QzRXpORWROUlRKa00xaDBOVFp2VjB4blNHeFBUalpWZDBObldVbExiMXBKQ25wcU1FVkJkMDFFWVZGQmQxcG5TWGhCVGxWRmNFaE5VMmhMYzJWTmIyUnVjR2h3WTFGdVUzWTBUWEJQTUZGNmNIa3JWelUzTm5KTVYydzRaRkJKYXpBS2EySnRWR3B3Y25NNU1XOWlVQzl5YjFoUlNYaEJUR3haSzJNMmJXOXlTMGx4YUVsR1ZGVnJlRzlTVEhsclRsaFRXVlpFV21vM1psWnlNMGMwYVRBeE5RcEVlRzlVTDBJM01EaE5Za1p3UnpWMVVWTk1iVWwzUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgqoIvpbI9JuzsEtCnPXA3Edri0qkxIV4I7KgV2WgRhYsCFQCO3jdTqRpnBYkiFdZ8zQcuCgcloRgPMjAyNjEwMTkxNjU5NDdaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE2NTk0N1owLwYJKoZIhvcNAQkEMSIEIDgdlDkadasHNL+udTK0McKrEGMFSnYj0meCm7Mk5SeQMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMHH/edAD7UOXy/ZApjV75KSYrqWamUSZHmrXt+QT3VNJah7UPyEoLPeP7g0XWKe38AIwe0TX/zdd51T5ZElxx+54X1ppgBtvEjx+6OesM2VLWG0ckT2o82Ki9uT2NkTLCR12"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"kYM5v9TC4mKEl1UHfkLeZOIsCZ7BP1EznAMJxuPHnqQ="}, "signature":"MEUCIQCLUcd2o0ZY6kfLaxqWKo/4eVth7hU1EpGwPIhRapV1uwIgLnb6e/hbmCynBVhzwED1Naj1d/I8rjcKH0LTdqgexIE="}} \ No newline at end of file diff --git a/build/torch214-cu126-x86_64-linux/symbols.json b/build/torch214-cu126-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch214-cu126-x86_64-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-x86_64-linux/__init__.py b/build/torch214-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch214-cu130-x86_64-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-x86_64-linux/_ops.py b/build/torch214-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch214-cu130-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch214-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..b6d5c68078fad4d0999a56998e5e344dd0dafa33 --- /dev/null +++ b/build/torch214-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bf8c7371171e1097e0c9bd75a3a9ee086b2462b0d4343b19ded87291d168a9bb +size 410568 diff --git a/build/torch214-cu130-x86_64-linux/metadata.json b/build/torch214-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..d2d4d1a51de73acc59a6101c15badd4f7b17dbe6 --- /dev/null +++ b/build/torch214-cu130-x86_64-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": "v4xzcRceEJfgyb11o6nuCGskYrDUNDsZ3thykdFoqbs=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch214-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..ba558f0b014e993bc6586ac652aed346517afd14 --- /dev/null +++ b/build/torch214-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLjCCBrOgAwIBAgIUdxcUA6MS5vnlQeGVzDYStJicoNQwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTY1OTUyWhcNMjYxMDE5MTcwOTUyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAESX83eLXGknxAsQEckkoeolKxSTh/AatkNnBVq3bvmg7Yk2nuusuUC1iAk1iEG9TqEcSsqwF7mQdeH/fkL0934qOCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRr0RwAABAMASDBGAiEAs3iVKZ8zd9TsvW1UYpsDDbpqrbUoOV0kZP+eDVM8ZawCIQCLUhAw3N53s/SThK7XnwgQgNiVm3RAWUCrPpRTck7dtzAKBggqhkjOPQQDAwNpADBmAjEArIUT6f80xdSAOgLNpXjMmi68JMIsOp6ljhEcxfbBe2GF/4w63skE+8zguNtXce6IAjEAkt1D0JHZf0NkUfGtvK0of9EpNQjs0YkHFGvSH6N1jeIsJLlO6FentDGkg4sXr7zt"}, "tlogEntries":[{"logIndex":"3323783884", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429195", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQDSsdmNBWrOMZGjv6JKEnfc3UdXekC4wkvDTvRgA3Bb3AIgPNaibd/k8p6clbRfs4nC4vZnG5kN0QKeiwfDLzhaM3I="}, "inclusionProof":{"logIndex":"3201879622", "rootHash":"WrucejyofTpSlYwa7Gg5RtHAuI3R/dngBh2xhseS1yc=", "treeSize":"3201880201", "hashes":["mQ8DGj/TGHhsjVCABK89Jr9LE9kcFDiSHuZcyb64Avc=", "sA8vmK/7pK6WFD/7ScaFlgsL0x1zFegHYG4Xkecyba4=", "4HjEiaNQiqBYmERNNKQ1OPDXeU7v/0c4Njn6XxsBhL8=", "X3UQ5cCBa+k798wiTM3agFe2meKc75pjTQ016ydKWNE=", "+SmKAXISHRaVW9Cp+HR7Akvo5Wx1CmNhANyxEqfUKOg=", "BNkXJThcvHBo9mQfDjH4FIh2KemDvgvOl2n7RCEfTUc=", "ZlyJCdbPpW1T/cValwXLLOur0HiKZ5D/rqtjhTc31sQ=", "J6Qh6UFwOstkg4xBSc+J40k1u2e02GxZdYnVaKFBEHA=", "67yY9qPQpeZdgvc058k9homghjVQVhTWtl/PU9qJa0c=", "XLpQF9pYuneKaqYHAZvYvquYx72QL+Orj1VTS/9j/ok=", "aakOu1RIwtqh91vCWaFfNB7xQile6QSDZODk6838XiU=", "PywsDAcGw1w2GA2yax7XMnye0HGAwlZxH8di0wVvUY8=", "rMd4uiSxqjI4gebUgztAJPb9ypgnt3Wjy/u01yzYF0E=", "qttRZVUdmA9Dz2AmZd4s0AFXdh4VB0rRnYnOMWDAB9o=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201880201\nWrucejyofTpSlYwa7Gg5RtHAuI3R/dngBh2xhseS1yc=\n\n— rekor.sigstore.dev wNI9ajBGAiEAlQCKdcw4LLhIZIp8vuP8ZWOppQORfACfaeIS6g56QhcCIQDhrjccjCEjUE8ERV2+N0quYmEe4ne5GO4cESisSZJ8xA==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJkZmI5ODZjZWMyNTg5OTQ3ZWViYzYyOGE2YWEyYTIxMjllYWJjNmJmNDA3YmVkYzQwYTg0ZjkwNDRhYmFhNzE3In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUURhZmk2QlR3Z1MrZzFmNDd5d1doSnhBbzYvOVphMG42dm1Pb3FuR21MRkJnSWhBS09KY2ZqZ0txTWNCemZ4UDZaTGFNOHgrQTgrTkpGMkppV3QvK1RiWXByViIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1ha05EUW5KUFowRjNTVUpCWjBsVlpIaGpWVUUyVFZNMWRtNXNVV1ZIVm5wRVdWTjBTbWxqYjA1UmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZGa3hUMVJWZVZkb1kwNU5hbGw0VFVSRk5VMVVZM2RQVkZWNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZUV0RnelpVeFlSMnR1ZUVGelVVVmphMnR2Wlc5c1MzaFRWR2d2UVdGMGEwNXVRbFlLY1ROaWRtMW5OMWxyTW01MWRYTjFWVU14YVVGck1XbEZSemxVY1VWalUzTnhkMFkzYlZGa1pVZ3ZabXRNTURrek5IRlBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKeU1GSjNRVUZDUVUxQlUwUkNSMEZwUlVGek0ybFdTMW80ZW1RNVZITjJWekZWV1hCelJFUmljSEZ5WWxWdlQxWXdhd3BhVUN0bFJGWk5PRnBoZDBOSlVVTk1WV2hCZHpOT05UTnpMMU5VYUVzM1dHNTNaMUZuVG1sV2JUTlNRVmRWUTNKUWNGSlVZMnMzWkhSNlFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2NFRkVRbTFCYWtWQmNrbFZWRFptT0RCNFpGTkJUMmRNVG5CWWFrMXRhVFk0U2sxSmMwOXdObXhxYUVWamVHWmlRbVV5UjBZS0x6UjNOak56YTBVck9IcG5kVTUwV0dObE5rbEJha1ZCYTNReFJEQktTRnBtTUU1clZXWkhkSFpMTUc5bU9VVndUbEZxY3pCWmEwaEdSM1pUU0RaT01RcHFaVWx6U2t4c1R6WkdaVzUwUkVkclp6UnpXSEkzZW5RS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgAcBiwSr0GEbnclhqwoGAiR+bSZKSBkNEJk2eML908eICFDNbYZIERkGghFnanZGY/BLuuRp2GA8yMDI2MTAxOTE2NTk1MlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTY1OTUyWjAvBgkqhkiG9w0BCQQxIgQg8o1lFsyFeBzvplc3m7JaSE2rJMLla/HlCtcVW4oC+7AwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAKj7+C3fZPjUO6kWAZODktMqu6dIKRrehoZSeZ2KFaFpSmqpQTfUClnjg4m0RAeILgIxAJo8IAtD1XLOEGMzjVuMBZGQ/ksJgJLOsqwIHg2aBCnXyzwUitq1KwUMDDgqNbB1sQ=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"37mGzsJYmUfuvGKKaqKiEp6rxr9Ae+3ECoT5BEq6pxc="}, "signature":"MEYCIQDafi6BTwgS+g1f47ywWhJxAo6/9Za0n6vmOoqnGmLFBgIhAKOJcfjgKqMcBzfxP6ZLaM8x+A8+NJF2JiWt/+TbYprV"}} \ No newline at end of file diff --git a/build/torch214-cu130-x86_64-linux/symbols.json b/build/torch214-cu130-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch214-cu130-x86_64-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-x86_64-linux/__init__.py b/build/torch214-cu132-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch214-cu132-x86_64-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-x86_64-linux/_ops.py b/build/torch214-cu132-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch214-cu132-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch214-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..faa463907e6d4947d64a2272ff1ea3a776f10092 --- /dev/null +++ b/build/torch214-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:15616061f42fae439f20412945c9bce9450ee01d2d0d39918e5b5d64251e73f1 +size 449416 diff --git a/build/torch214-cu132-x86_64-linux/metadata.json b/build/torch214-cu132-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..81acbe8bbfe85f0a9a8e9a160284feec643a8e94 --- /dev/null +++ b/build/torch214-cu132-x86_64-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": "FWFgYfQvrkOfIEEpRcm86UUO4B0tDTmRjltdZCUec/E=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch214-cu132-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..136948e7a668f5f30f8f1f1fc7893a1bfc6d5d95 --- /dev/null +++ b/build/torch214-cu132-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrGgAwIBAgIUTMpSuf7fF3GGiBuLbAJuVSOIbSMwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTcwMDAxWhcNMjYxMDE5MTcxMDAxWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEKuG/aPQD6wAwMAxcgWVhZOm63PtjJ/gstIGhjFw8EaZWtue6vxhKmfo47e1RJR4WWwk1j+xQ45wbggdi1ccjvqOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRsZUQAABAMARjBEAiAU2kVf7KTA3jeM2BRfjEfgsu8QsrinLeE2qDrrq6JF9AIge/7d80UyinCXr8InIeq0LFOWXZi5DCygTNa68SYpYvAwCgYIKoZIzj0EAwMDaQAwZgIxAK5xvsjn4V1md0in/TvXPeQ1cAuURk9kjzDRszvNqLzA3URojsUyKwiUdMADzRUEJAIxAO6BsmX0lduZ60ednOlFbh7efQAHNfB5nHaaidiyE414FnFmmjq0iSaKvMJ8b3tjVA=="}, "tlogEntries":[{"logIndex":"3323785233", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429202", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQD4Ka14z6nk9a+98SE6ueuz9kxB+qbJSHJyej0fK0J48AIgTJE0c6bZowQiWNKYddvVj4enLkXUHqXdTC9oT9Ykdno="}, "inclusionProof":{"logIndex":"3201880971", "rootHash":"VY9F1pXM0/vsg4OO0IL88Uf8O/nRRrpo77ZBil1iC0g=", "treeSize":"3201881246", "hashes":["Rt/HSwP2iP5HsC18O9xw71nwLDIboC8V4W/lENmICJ8=", "MCu1uvEL4gUz0Qj08aJWWTMKXg1/A0IWgPFtSd5Q9Yg=", "nig3Qq9wBv2O7m+++r1h2zV3EQmIJX0sHWI82i7nbHo=", "eUJByEkWFfjTq07lEwEgaCcxVMiUXmrRFoIGwgH3Cek=", "66Vum0S5yTedgSh9Hi5EMiF5bUGWGaRPaXr5Zjp7GV0=", "qXDY/C+ZCb6vGZg8HdOk6ons2nqHgNLl/bs2++JgJcY=", "e6xDVSWou/CFSgPsZ/OE/G4wtHdRz+c2+ICvUEBwG4w=", "oyf61RvKy81mzSlAKlbbv9wiNOMGOty3zfNXQj5lMvc=", "wRhxcRTuxz7gqbsK5b6p9UX/yHkjhy3zkQYtSX3/YUU=", "ZrbCD2QlP/FLb4aSXmzXl+RPtVP2i/I7ZbyjhFm7cGk=", "1j8g/pYFV78pmfbAESBOYXCLMamXmuspYGnaTMfz4UE=", "jmu3TtaRxatojZcOxXho2Bka8bsanj6tbxXarHVTiec=", "qttRZVUdmA9Dz2AmZd4s0AFXdh4VB0rRnYnOMWDAB9o=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201881246\nVY9F1pXM0/vsg4OO0IL88Uf8O/nRRrpo77ZBil1iC0g=\n\n— rekor.sigstore.dev wNI9ajBFAiBrhBAd89Pr2ri1EVMRGZAgWV7p/DaZ1rrJ4VOBkaZIYQIhAPWP9y8sBmOvHQWiC39YxZnZKqx0qdRi2aeDkQCEoMjt\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI2OGFjNzI2YjA0MWViYTkxNDNhYjZlMzc0NDU3Y2RmOTcyNWFjYjlkZDRkODY2Mjg1N2Q1MDY3NTVjZTUxYTIyIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNtaGQ2bmNocU1hTTI0MU9scjhBTHQ3cHBOcXFoU1NHam9SNit5b0dlT1ZnSWdlRzlnTXlRM0o2RVJWMWdNQjJLdGp2cjFLSUptV2lPcHdDMXJoWHpxUXc0PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KSFowRjNTVUpCWjBsVlZFMXdVM1ZtTjJaR00wZEhhVUoxVEdKQlNuVldVMDlKWWxOTmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZHTjNUVVJCZUZkb1kwNU5hbGw0VFVSRk5VMVVZM2hOUkVGNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZMZFVjdllWQlJSRFozUVhkTlFYaGpaMWRXYUZwUGJUWXpVSFJxU2k5bmMzUkpSMmdLYWtaM09FVmhXbGQwZFdVMmRuaG9TMjFtYnpRM1pURlNTbEkwVjFkM2F6RnFLM2hSTkRWM1ltZG5aR2t4WTJOcWRuRlBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKeldsVlJRVUZDUVUxQlVtcENSVUZwUVZVeWExWm1OMHRVUVROcVpVMHlRbEptYWtWbVozTjFPRkZ6Y21sdVRHVkZNZ3B4UkhKeWNUWktSamxCU1dkbEx6ZGtPREJWZVdsdVExaHlPRWx1U1dWeE1FeEdUMWRZV21rMVJFTjVaMVJPWVRZNFUxbHdXWFpCZDBObldVbExiMXBKQ25wcU1FVkJkMDFFWVZGQmQxcG5TWGhCU3pWNGRuTnFialJXTVcxa01HbHVMMVIyV0ZCbFVURmpRWFZWVW1zNWEycDZSRkp6ZW5aT2NVeDZRVE5WVW04S2FuTlZlVXQzYVZWa1RVRkVlbEpWUlVwQlNYaEJUelpDYzIxWU1HeGtkVm8yTUdWa2JrOXNSbUpvTjJWbVVVRklUbVpDTlc1SVlXRnBaR2w1UlRReE5BcEdia1p0YldweE1HbFRZVXQyVFVvNFlqTjBhbFpCUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgqzzu+MX23MpK3GEo3AnkjTIdHcEB4cvyRovNM4LO+vgCFQDnOTRFFGIwVvWXNXtsMqLTYm8I1xgPMjAyNjEwMTkxNzAwMDFaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE3MDAwMVowLwYJKoZIhvcNAQkEMSIEIDc29TEGSlQ7KSSGdQPVWYW5itNS8rIRA5vbTDmfXsnnMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQDAX3zg+q7g8c42p/Jns9MoYPcZqiU98TuE4DD0mfESGNJ4f6mrhZGi5agkmITzPkACMQDgIqxLxI6SEFJHFFM1xKlJJuZuftDe86n9p5FGb265X29wn6FXKoiaVI8Zsun8SZ0="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"aKxyawQeupFDq243RFfN+XJay53U2GYoV9UGdVzlGiI="}, "signature":"MEUCIQCmhd6nchqMaM241Olr8ALt7ppNqqhSSGjoR6+yoGeOVgIgeG9gMyQ3J6ERV1gMB2Ktjvr1KIJmWiOpwC1rhXzqQw4="}} \ No newline at end of file diff --git a/build/torch214-cu132-x86_64-linux/symbols.json b/build/torch214-cu132-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch214-cu132-x86_64-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-x86_64-linux/__init__.py b/build/torch215-cu130-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch215-cu130-x86_64-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-x86_64-linux/_ops.py b/build/torch215-cu130-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch215-cu130-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch215-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..e07196f312f15c392764a0016c0556aea59f804b --- /dev/null +++ b/build/torch215-cu130-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4ddedbe19ff7ae14e00ec1c06059eee696419b18ddd260196cfa44297c21d82c +size 410568 diff --git a/build/torch215-cu130-x86_64-linux/metadata.json b/build/torch215-cu130-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..613e4f62f6fa84b87c288fd7e1fef2f0c6cd8664 --- /dev/null +++ b/build/torch215-cu130-x86_64-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": "Td7b4Z/3rhTgDsHAYFnu5pZBmxjd0mAZbPpEKXwh2Cw=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch215-cu130-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..d0cc2fc69df7769ac1b0be49395d0679c47cc10c --- /dev/null +++ b/build/torch215-cu130-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrKgAwIBAgIUM1mO0eYFE+nwrK0twXz6Mih3Y2IwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTcwMDA2WhcNMjYxMDE5MTcxMDA2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEUsuCRwA8PKEWfO0gYOisoZbEDq4mOhen5DhwjA6Z8DXD5WHJZnMvG2/wCcaF4NjQvwNtOUnVlCiXKiWsELf7mqOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRsq4wAABAMARzBFAiAWp4lw4mSjmqDItwYlCJOtoFWv374QAQFulmQ3Q9HJ+wIhALM89kQ53MWb9kR/MWKxuG+2aBMRbIuK6ZwAYi0cRvLGMAoGCCqGSM49BAMDA2kAMGYCMQDTEPP+owhnBKDv3yHo6HhlRj8bOsPCpIUR5Qv/n+2rVGt12EIOI54LJp8iLJjAel8CMQCsTCAJ7w4WfF7CL3UhFhNb/jauE1MuYX1/3qQOQt4MkQycL2zoPey2Alz9I6LdVeY="}, "tlogEntries":[{"logIndex":"3323785940", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429206", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQC+q+0HFcnUdLzCg8/4sL59kQAlPJz4tMHGI3M1MfK+3gIgT+UP6ArcruzhAJtwFuRUpOhNeZ2qP/qgbnpBaeflKxY="}, "inclusionProof":{"logIndex":"3201881678", "rootHash":"jDAp0AeaP96Ylg40SIGBenZ33EqEnlA7YmlLxxGp3VM=", "treeSize":"3201881722", "hashes":["yX15AlIUl8iaG3tz1HL4g9GJ7roH4tjYvyydPus7+dw=", "id8+9+Rb4NvxMIAHkVaPTTFtQJ03JcqpkYnGwRM8Z/k=", "8GFoRq0bMZkKfr+utbsew5lIfgL3SrphfBRejk67V64=", "UcrEg8Hi4kVQyQMQJfwgOgrPo6qrpsWAfkwg1EH19M0=", "yIMq5diU/wH0TQZSCF/m42FrI3f/kikaK3VpEjW8m/g=", "K38YeK8j4RWtS4xbmZrgENLkV6GO3M2M81MzOm5J8N0=", "k3sO5wQJIlc2yQY17EbOYMp2VOGfxLMo9BeuyoA+944=", "WLgHeG86OdsrhffSSkr63B+o0hOL/ZovwVe7b5FviFU=", "jQYKjCqOEppZ3aOYixOwy8Ev0SZ3ZDsanM++4zKVFjc=", "jmu3TtaRxatojZcOxXho2Bka8bsanj6tbxXarHVTiec=", "qttRZVUdmA9Dz2AmZd4s0AFXdh4VB0rRnYnOMWDAB9o=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201881722\njDAp0AeaP96Ylg40SIGBenZ33EqEnlA7YmlLxxGp3VM=\n\n— rekor.sigstore.dev wNI9ajBEAiBeCXAKBii4lOAdnoiho1GHIifR4u5B1vSeB0NRjP/BSQIgEAeX3F0YllOF2RLeLOSU1vbYHkTw+e3UFjAu6aytGtY=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI5NGI3OWJhZTNjZWIzMmNjNWJhYjcwYTUwNjA5NDdhOGM4YTc5M2ZiMWQ0MDY5NjEwNzljMTQyNzhlODEzYzViIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUQ4MkhHM2NYZUEwdmJMbWM0SlMvc0tET3FLQnduMU1vVzd2ZVVrMWphT253SWhBS3dLdWg5SmdnT2k0Wi9kUVc5VUVubGx2Y0hxbnh3L2lkRXY0Wk0zNUJoRiIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KTFowRjNTVUpCWjBsVlRURnRUekJsV1VaRksyNTNja3N3ZEhkWWVqWk5hV2d6V1RKSmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZHTjNUVVJCTWxkb1kwNU5hbGw0VFVSRk5VMVVZM2hOUkVFeVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZWYzNWRFVuZEJPRkJMUlZkbVR6Qm5XVTlwYzI5YVlrVkVjVFJ0VDJobGJqVkVhSGNLYWtFMldqaEVXRVExVjBoS1dtNU5ka2N5TDNkRFkyRkdORTVxVVhaM1RuUlBWVzVXYkVOcFdFdHBWM05GVEdZM2JYRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKemNUUjNRVUZDUVUxQlVucENSa0ZwUVZkd05HeDNORzFUYW0xeFJFbDBkMWxzUTBwUGRHOUdWM1l6TnpSUlFWRkdkUXBzYlZFelVUbElTaXQzU1doQlRFMDRPV3RSTlROTlYySTVhMUl2VFZkTGVIVkhLekpoUWsxU1lrbDFTelphZDBGWmFUQmpVblpNUjAxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKclFVMUhXVU5OVVVSVVJWQlFLMjkzYUc1Q1MwUjJNM2xJYnpaSWFHeFNhamhpVDNOUVEzQkpWVkkxVVhZdmJpc3ljbFpIZERFS01rVkpUMGsxTkV4S2NEaHBURXBxUVdWc09FTk5VVU56VkVOQlNqZDNORmRtUmpkRFRETlZhRVpvVG1JdmFtRjFSVEZOZFZsWU1TOHpjVkZQVVhRMFRRcHJVWGxqVERKNmIxQmxlVEpCYkhvNVNUWk1aRlpsV1QwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg6qBwYp9YSJSBoyghPn0aeXtadC77Es0ebzsptcHudPgCFQCCL4JRQX/oS/N8qjRJntCeMOOUPhgPMjAyNjEwMTkxNzAwMDZaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE3MDAwNlowLwYJKoZIhvcNAQkEMSIEIDPWUNPWPpDa6WaaJQQkHsHWi1EteRsDfz1K5vkiTRa6MIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMC3qz9TCYdErUoQKWnmry1nVj1gZYHGHDwdYD4edmUUr4yyVNtVcOGZI1QcDtNRyHAIxALxI/6Sis0j8OsAqlbOzZCcDi+HQrNz9Idyi4TcsAaUdpZN56f5t3MXF2gQSeEA+Og=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"lLebrjzrMsxbq3ClBglHqMink/sdQGlhB5wUJ46BPFs="}, "signature":"MEYCIQD82HG3cXeA0vbLmc4JS/sKDOqKBwn1MoW7veUk1jaOnwIhAKwKuh9JggOi4Z/dQW9UEnllvcHqnxw/idEv4ZM35BhF"}} \ No newline at end of file diff --git a/build/torch215-cu130-x86_64-linux/symbols.json b/build/torch215-cu130-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch215-cu130-x86_64-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-x86_64-linux/__init__.py b/build/torch215-cu132-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch215-cu132-x86_64-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-x86_64-linux/_ops.py b/build/torch215-cu132-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch215-cu132-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch215-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..023461ec89bef90030a276921ef5cb317d5da794 --- /dev/null +++ b/build/torch215-cu132-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cb2d25087e6ed4fc6954bf5f502884e7fec32fc76e318c2346b6985d2aec53e9 +size 449416 diff --git a/build/torch215-cu132-x86_64-linux/metadata.json b/build/torch215-cu132-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..557c87c644c5944a969c74c6b4f8353c0650d58f --- /dev/null +++ b/build/torch215-cu132-x86_64-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": "yy0lCH5u1PxpVL9fUCiE5/7DL8duMYwjRraYXSrsU+k=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch215-cu132-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..c39893adf68ddef959f812a0d2d562ebfb11d0ec --- /dev/null +++ b/build/torch215-cu132-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrGgAwIBAgIUMps4I5kcWR0VRbdSHpMr+/eR/ugwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTcwMDA3WhcNMjYxMDE5MTcxMDA3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAELlXddG9lL15mjp8iJdNrRdhjHGsvLIDEv6ZZr7E+7i1gxhzmVIdj3mhPAHytX9JBQ+AP1sc0m3ybNmMaBugzXKOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRsvIQAABAMARjBEAiBJeQzppIBSjZvTQKl52OruQwtBwXm/yuhH/HbbTfzT8gIgKINlQYOOAqEn+hTMuCL1X3o0IGNG0IfxeXVBQo9Vh9wwCgYIKoZIzj0EAwMDaQAwZgIxAO0lgSSmfli7mn7/ap3T/vJNbojUZE5G60LGBKKCtSxIq+fP2yGqr/ZEAm5nR7DQFQIxAPWR9pYRT2uPIp0IytwS9snky9enR+EoW+wSjRmmmCVS2oE7uVdH5ubA03ikMZQOIQ=="}, "tlogEntries":[{"logIndex":"3323786158", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429207", "inclusionPromise":{"signedEntryTimestamp":"MEUCIFy33q89lYHYdQL6BJcNXiIo7alul/eNwh+KGPfJ0hzSAiEA9WAiipAYwt3lzovXiKOE0DBvgnNjeEVcDXguOFl9hhE="}, "inclusionProof":{"logIndex":"3201881896", "rootHash":"SlJN76BbURUYrxQrXmyXn5UqdRaZAhS5Fa5XuzxR3dU=", "treeSize":"3201881904", "hashes":["6wzSDdf38iynUa/WLI2uU95Izh48oUvr55zeO+9wtO0=", "xb6paBoHKi9zxXPJfG3/Q5XHpAUyNiYrP1T7ufLbuhI=", "3Q9+G4QyF5OzgvZxxoyIpjz5S9SIgiu4NSTxNwaPDiI=", "Tj1/E/YjxnISUC5GfdsfIqUPZgWUveDyTiATRv7IWHM=", "CLnimObtuDaDSR3iN+NRGRMOul5RYbVKn9+1BwAiajY=", "P211B5XN7jv23SZ3FBLAM3W25ErHu+5nTFY8lMRJsPQ=", "WLgHeG86OdsrhffSSkr63B+o0hOL/ZovwVe7b5FviFU=", "jQYKjCqOEppZ3aOYixOwy8Ev0SZ3ZDsanM++4zKVFjc=", "jmu3TtaRxatojZcOxXho2Bka8bsanj6tbxXarHVTiec=", "qttRZVUdmA9Dz2AmZd4s0AFXdh4VB0rRnYnOMWDAB9o=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201881904\nSlJN76BbURUYrxQrXmyXn5UqdRaZAhS5Fa5XuzxR3dU=\n\n— rekor.sigstore.dev wNI9ajBFAiEAzbBhGhObHcdf8jfZAfggUHVsgzdM05G9hBoxIr5CGpcCIAHJp/TWNDJ0fGaOrS9jfrXnqtbxJbw6rKiTLKekiMNr\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3NjQ3NDVjNDBlMGU2MjUzNGVmNTA5MTFmZTA3ZmYyNjcxMzhhOTJhMmE1NTlhMTczMzdjZTkyMzc4YjhkZTZlIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRGNJZWk1STZlbytsaVZmekY0cXNoQzBHUkF3Yy8ra1BJUytQR3BKTzBHOUFpRUF1Y1dHWkZwZDd3djc1YW12UXBYbVgzM0NnekpJd2s4UEJDQThLT1c5blVVPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KSFowRjNTVUpCWjBsVlRYQnpORWsxYTJOWFVqQldVbUprVTBod1RYSXJMMlZTTDNWbmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZHTjNUVVJCTTFkb1kwNU5hbGw0VFVSRk5VMVVZM2hOUkVFelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZNYkZoa1pFYzViRXd4TlcxcWNEaHBTbVJPY2xKa2FHcElSM04yVEVsRVJYWTJXbG9LY2pkRkt6ZHBNV2Q0YUhwdFZrbGthak50YUZCQlNIbDBXRGxLUWxFclFWQXhjMk13YlRONVlrNXRUV0ZDZFdkNldFdFBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKemRrbFJRVUZDUVUxQlVtcENSVUZwUWtwbFVYcHdjRWxDVTJwYWRsUlJTMncxTWs5eWRWRjNkRUozV0cwdmVYVm9TQW92U0dKaVZHWjZWRGhuU1dkTFNVNXNVVmxQVDBGeFJXNHJhRlJOZFVOTU1WZ3piekJKUjA1SE1FbG1lR1ZZVmtKUmJ6bFdhRGwzZDBObldVbExiMXBKQ25wcU1FVkJkMDFFWVZGQmQxcG5TWGhCVHpCc1oxTlRiV1pzYVRkdGJqY3ZZWEF6VkM5MlNrNWliMnBWV2tVMVJ6WXdURWRDUzB0RGRGTjRTWEVyWmxBS01ubEhjWEl2V2tWQmJUVnVVamRFVVVaUlNYaEJVRmRTT1hCWlVsUXlkVkJKY0RCSmVYUjNVemx6Ym10NU9XVnVVaXRGYjFjcmQxTnFVbTF0YlVOV1V3b3liMFUzZFZaa1NEVjFZa0V3TTJsclRWcFJUMGxSUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgGZz0teMcoGSV9SSrfX4l7ow3m1UMVAAyJnglLKuFiHICFQDPNfFlOPWaSoAfL1R1U8DXEgW9AhgPMjAyNjEwMTkxNzAwMDdaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHaMIIB1gIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE3MDAwN1owLwYJKoZIhvcNAQkEMSIEIOlU3lAJF27iBmw0RoW+Wds006Aa1Xkd7H1wtO3fgFyjMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRmMGQCMHQ638tdRHSeeRKr3e+C4UDsKgFRZcIjLW3DGxd+7zFddkYC5APqAXeCCDGmZwxEuAIwIoqyNbrMoSiplzUVolwsk31kjCxYzuhZqhM0k8KqUF1nrJOzmy30N/WtsA4dI9se"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"dkdFxA4OYlNO9QkR/gf/JnE4qSoqVZoXM3zpI3i43m4="}, "signature":"MEUCIDcIei5I6eo+liVfzF4qshC0GRAwc/+kPIS+PGpJO0G9AiEAucWGZFpd7wv75amvQpXmX33CgzJIwk8PBCA8KOW9nUU="}} \ No newline at end of file diff --git a/build/torch215-cu132-x86_64-linux/symbols.json b/build/torch215-cu132-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch215-cu132-x86_64-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-x86_64-linux/__init__.py b/build/torch215-cu134-x86_64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea22dbc160d2b75b32c3c8c5d882d25c77b494 --- /dev/null +++ b/build/torch215-cu134-x86_64-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-x86_64-linux/_ops.py b/build/torch215-cu134-x86_64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..91b9d7d1e1000d230a967add3225cb508a07fe99 --- /dev/null +++ b/build/torch215-cu134-x86_64-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-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so b/build/torch215-cu134-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..983d5f6c602bc7f13edf3aa2abe0dc1e8479d17b --- /dev/null +++ b/build/torch215-cu134-x86_64-linux/_rwkv_cuda_fe2bb66.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:37d2af48b312b066d6e4e7b166d7a75c837683dd24153378f4088e012729b350 +size 428024 diff --git a/build/torch215-cu134-x86_64-linux/metadata.json b/build/torch215-cu134-x86_64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..2249a7700fe991879f0442b0dc3f22985ab07bb4 --- /dev/null +++ b/build/torch215-cu134-x86_64-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": "N9KvSLMSsGbW5OexZtenXIN2g90kFTN49AiOAScps1A=", + "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-x86_64-linux/metadata.json.sigstore b/build/torch215-cu134-x86_64-linux/metadata.json.sigstore new file mode 100644 index 0000000000000000000000000000000000000000..03729bb1e31ba81e776d2a4b537e5773e4e4f13d --- /dev/null +++ b/build/torch215-cu134-x86_64-linux/metadata.json.sigstore @@ -0,0 +1 @@ +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrKgAwIBAgIUal/87Pl9MOJyB+3dXZehw6JEzJMwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTcwMDExWhcNMjYxMDE5MTcxMDExWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAET/4p+enLBbr2MAfFAWagZr0MiauU6rOCe4FSjiEiofU7PIyvT5lYBV9FRPi/PuX3c7w2BtmXN8rGMaCJlKEIBaOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTY1MTI0OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhVRs/gAAABAMARzBFAiA4fxNu8lFJfdxlnDkhgEbGSJnxX2bgxUu99j3/l2l7ZwIhAMhzOgOqbompxOUGEhbxyTAhXgtorS7EebCajmJlRwcpMAoGCCqGSM49BAMDA2gAMGUCMG+c17qNJKlVr+QzTqHu1appmlcJSJyJI8K7DxK44KoTtWS5zdcEqtbjlsNwPQuHwwIxAMfiKSh/xsnBL66fgjK25xg5fBwNNnaGA6TJQfwowZHzsYMB8fV+PiXmYmMr9GJuQA=="}, "tlogEntries":[{"logIndex":"3323786978", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792429211", "inclusionPromise":{"signedEntryTimestamp":"MEUCIFQ2h6Dio+uwhbqjBTrAibH+4uK5935eQ4IIbpcHtr66AiEAkxh4BRO/r4rnsj+v847qwyawzBwPHcIUsUBYuDeBPz0="}, "inclusionProof":{"logIndex":"3201882716", "rootHash":"dyMmWbxhAi8VB+lAAo2nGxFLmInmEazb3VQPddlMRbc=", "treeSize":"3201882962", "hashes":["jbQI2m/9+LeM9QnfgFxNcizGkfw+wgUjsUZjIkQgyRY=", "65dsGYu/fV5cPJ4wioXfbRqi+xqfBnxV0SXnEqqqE/w=", "WtlrKW2k2hJ1roC6XB3nr7DGowToIjLIIW3URgpxywE=", "imrtrF3olE9a/1wjQtauZwxRXp1p1jmaRDkX6CfKjRw=", "H863us9YXxJtLURnybRyfAigmTXZFjzPcYk4lXKXOtE=", "PTu6fly2TBj8Wb56SHyr09Q4TK6Tfu+ZSZv9Xaxshi8=", "rsV77MyexAb/GAGhrAV9fBwCIGXcv8si8R4G10e7qrM=", "24WhL7k1ubhXy36etA2ER84unlCU/mTDBhiQAKvcfZ4=", "KrVVvohLrx+SB0eS1Noz9gmSPC+wKX3rmKLRpI/3FRs=", "XgVjE9eQeGwefPmVzFSwzA1ET8swDM8PTndYtYz8yT8=", "y8nV9QA2dN6jLyjE1ImPSkI5nVAVa/xhHfZsDxInkeg=", "jmu3TtaRxatojZcOxXho2Bka8bsanj6tbxXarHVTiec=", "qttRZVUdmA9Dz2AmZd4s0AFXdh4VB0rRnYnOMWDAB9o=", "8EBu2br9+uQVSURl2FbPBD5EU28cYYhiRWnWgduwE4I=", "MD4OgR/R/G+2HawrFMhrsxg9GvdgOeQsIipPmrKXl38=", "j0WYhlZCgHWFR/OJIjCQT0Q8ta+HnlghbtUUUTyZptI=", "7WFpoaiQCzLwM3o2YI5Ngo8J7JD0OfyCCQtmXhWxVes=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3201882962\ndyMmWbxhAi8VB+lAAo2nGxFLmInmEazb3VQPddlMRbc=\n\n— rekor.sigstore.dev wNI9ajBEAiBJ/F7UUjV1H0Xyk/WUTvpLSuOCOTq6BojrwtOpewt70gIgAdtjGrIw4KMZAqJqCQ9LYKG849+6USyrPBOB/OQMHNI=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIzZTliZTNlZDFiMWM4Mzk4NjkzODU3NDViN2FhNTRlMTUxZDQwYTVmYTNhNzhlYjc1NzQ4YTFjYjIzNjA4YzZkIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJRVJZU0NINnBLcHhmeFBiRnZZWDJrZnRVaWJwWkROa1RFeGJnd20yRUJydkFpQUdxUHpkL3RNZUF4dTRsWFJabm5MT05MRU81M0RZbnI3YzRFSFlla2Z3dVE9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KTFowRjNTVUpCWjBsVllXd3ZPRGRRYkRsTlQwcDVRaXN6WkZoYVpXaDNOa3BGZWtwTmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZHTjNUVVJGZUZkb1kwNU5hbGw0VFVSRk5VMVVZM2hOUkVWNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZVTHpSd0syVnVURUppY2pKTlFXWkdRVmRoWjFweU1FMXBZWFZWTm5KUFEyVTBSbE1LYW1sRmFXOW1WVGRRU1hsMlZEVnNXVUpXT1VaU1VHa3ZVSFZZTTJNM2R6SkNkRzFZVGpoeVIwMWhRMHBzUzBWSlFtRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVV1RGTlZFa3dUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVmxKekwyZEJRVUZDUVUxQlVucENSa0ZwUVRSbWVFNTFPR3hHU21aa2VHeHVSR3RvWjBWaVIxTktibmhZTW1KbmVGVjFPUW81YWpNdmJESnNOMXAzU1doQlRXaDZUMmRQY1dKdmJYQjRUMVZIUldoaWVIbFVRV2hZWjNSdmNsTTNSV1ZpUTJGcWJVcHNVbmRqY0UxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKblFVMUhWVU5OUnl0ak1UZHhUa3BMYkZaeUsxRjZWSEZJZFRGaGNIQnRiR05LVTBwNVNrazRTemRFZUVzME5FdHZWSFJYVXpVS2VtUmpSWEYwWW1wc2MwNTNVRkYxU0hkM1NYaEJUV1pwUzFOb0wzaHpia0pNTmpabVoycExNalY0WnpWbVFuZE9UbTVoUjBFMlZFcFJabmR2ZDFwSWVncHpXVTFDT0daV0sxQnBXRzFaYlUxeU9VZEtkVkZCUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgFn7PE4gwboFKJigOgKoAdHDh7FZEOVYLXOtcxiVShqcCFQCjCT4Q1rrzNJiGLqnpp+BQf6rlixgPMjAyNjEwMTkxNzAwMTFaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAxOTE3MDAxMVowLwYJKoZIhvcNAQkEMSIEIL9+GYfsPmD5KrgxcJuMQICadgJN4PIhGK8H18jx3ULkMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMBaOOfihv60XArohldJ90eIvq19xamIqHR6+uDSAw5DCZxl3RxbD0F5dqCtoP1mIWAIxAPs5UkVX2juTzT/EnoriIa8UeOBBSH2w6OfrORw0AlMvOAtlKmbxVumJmRtyTgShiQ=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"Ppvj7Rscg5hpOFdFt6pU4VHUCl+jp463V0ihyyNgjG0="}, "signature":"MEQCIERYSCH6pKpxfxPbFvYX2kftUibpZDNkTExbgwm2EBrvAiAGqPzd/tMeAxu4lXRZnnLONLEO53DYnr7c4EHYekfwuQ=="}} \ No newline at end of file diff --git a/build/torch215-cu134-x86_64-linux/symbols.json b/build/torch215-cu134-x86_64-linux/symbols.json new file mode 100644 index 0000000000000000000000000000000000000000..a917662559ad66a0af323fb25aa4545fddddd70e --- /dev/null +++ b/build/torch215-cu134-x86_64-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": [] +}