Uploaded using `kernel-builder`.

Browse files

Files changed (16) hide show

build/torch212-cxx11-cu130-x86_64-linux/{_causal_conv1d_state_cuda_22d776b.abi3.so → _causal_conv1d_state_cuda_8589753.abi3.so} +1 -1
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu130-x86_64-linux/causal_conv1d_state/__init__.py +0 -26
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +21 -2
build/torch212-cxx11-cu132-x86_64-linux/{_causal_conv1d_state_cuda_22d776b.abi3.so → _causal_conv1d_state_cuda_8589753.abi3.so} +1 -1
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu132-x86_64-linux/causal_conv1d_state/__init__.py +0 -26
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +21 -2
build/torch213-cxx11-cu130-x86_64-linux/__init__.py +305 -0
build/torch213-cxx11-cu130-x86_64-linux/_causal_conv1d_state_cuda_8589753.abi3.so +3 -0
build/torch213-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cu130-x86_64-linux/metadata.json +41 -0
build/torch213-cxx11-cu132-x86_64-linux/__init__.py +305 -0
build/torch213-cxx11-cu132-x86_64-linux/_causal_conv1d_state_cuda_8589753.abi3.so +3 -0
build/torch213-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cu132-x86_64-linux/metadata.json +41 -0

build/torch212-cxx11-cu130-x86_64-linux/{_causal_conv1d_state_cuda_22d776b.abi3.so → _causal_conv1d_state_cuda_8589753.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8f1362645c438c155acbc4add3a0b8dc8dd3e1e677f1bc20030e13739b879fba
 size 1256392

 version https://git-lfs.github.com/spec/v1
+oid sha256:3b22d80b41c587dbf1b30addb26be2c65c37168871b823b26548d8c8109053b1
 size 1256392

build/torch212-cxx11-cu130-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _causal_conv1d_state_cuda_22d776b
-ops = torch.ops._causal_conv1d_state_cuda_22d776b
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_causal_conv1d_state_cuda_22d776b::{op_name}"

 import torch
+from . import _causal_conv1d_state_cuda_8589753
+ops = torch.ops._causal_conv1d_state_cuda_8589753
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_causal_conv1d_state_cuda_8589753::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/causal_conv1d_state/__init__.py DELETED Viewed

@@ -1,26 +0,0 @@
-import ctypes
-import importlib.util
-import sys
-from pathlib import Path
-from types import ModuleType
-def _import_from_path(file_path: Path) -> ModuleType:
-    # We cannot use the module name as-is, after adding it to `sys.modules`,
-    # it would also be used for other imports. So, we make a module name that
-    # depends on the path for it to be unique using the hex-encoded hash of
-    # the path.
-    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
-    module_name = path_hash
-    spec = importlib.util.spec_from_file_location(module_name, file_path)
-    if spec is None:
-        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
-    module = importlib.util.module_from_spec(spec)
-    if module is None:
-        raise ImportError(f"Cannot load module {module_name} from spec")
-    sys.modules[module_name] = module
-    spec.loader.exec_module(module)  # type: ignore
-    return module
-globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cu130-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "causal-conv1d-state",
-  "id": "_causal_conv1d_state_cuda_22d776b",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -18,5 +18,24 @@
       "8.9",
       "9.0"
     ]
   }
-}

 {
   "name": "causal-conv1d-state",
+  "id": "_causal_conv1d_state_cuda_8589753",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],
       "8.9",
       "9.0"
     ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "jC3cFEzRQxoloALJMANJU4Px3BpryCaxyVp+MdmTj6Q=",
+      "_causal_conv1d_state_cuda_8589753.abi3.so": "OyLYC0HFh9vxswrdsmvixlw3FohxuCOyZUjYyBCQU7E=",
+      "_ops.py": "XOwQje24M0e2Is75hhdYmSaw60wjwb/z4WY2Fza5UIc="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "81f55ea30fd8f819dcf93a3c934dd584c895bd2f",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "85897539ddf848f1cbd4bb17d24c89030d2c33e2",
+      "dirty": false
+    }
   }
+}

build/torch212-cxx11-cu132-x86_64-linux/{_causal_conv1d_state_cuda_22d776b.abi3.so → _causal_conv1d_state_cuda_8589753.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a06be13f796b15cacf2bae099c1b4b4c4aeb668c8cb0b5362172c8f87d9c0652
 size 1264584

 version https://git-lfs.github.com/spec/v1
+oid sha256:9a19590b37857c264d674fd768c1a98fdc74cd4def6234f7d7883575bcac9b13
 size 1264584

build/torch212-cxx11-cu132-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _causal_conv1d_state_cuda_22d776b
-ops = torch.ops._causal_conv1d_state_cuda_22d776b
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_causal_conv1d_state_cuda_22d776b::{op_name}"

 import torch
+from . import _causal_conv1d_state_cuda_8589753
+ops = torch.ops._causal_conv1d_state_cuda_8589753
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_causal_conv1d_state_cuda_8589753::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/causal_conv1d_state/__init__.py DELETED Viewed

@@ -1,26 +0,0 @@
-import ctypes
-import importlib.util
-import sys
-from pathlib import Path
-from types import ModuleType
-def _import_from_path(file_path: Path) -> ModuleType:
-    # We cannot use the module name as-is, after adding it to `sys.modules`,
-    # it would also be used for other imports. So, we make a module name that
-    # depends on the path for it to be unique using the hex-encoded hash of
-    # the path.
-    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
-    module_name = path_hash
-    spec = importlib.util.spec_from_file_location(module_name, file_path)
-    if spec is None:
-        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
-    module = importlib.util.module_from_spec(spec)
-    if module is None:
-        raise ImportError(f"Cannot load module {module_name} from spec")
-    sys.modules[module_name] = module
-    spec.loader.exec_module(module)  # type: ignore
-    return module
-globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cu132-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "causal-conv1d-state",
-  "id": "_causal_conv1d_state_cuda_22d776b",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -18,5 +18,24 @@
       "8.9",
       "9.0"
     ]
   }
-}

 {
   "name": "causal-conv1d-state",
+  "id": "_causal_conv1d_state_cuda_8589753",
   "version": 2,
   "license": "Apache-2.0",
   "python-depends": [],
       "8.9",
       "9.0"
     ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "jC3cFEzRQxoloALJMANJU4Px3BpryCaxyVp+MdmTj6Q=",
+      "_causal_conv1d_state_cuda_8589753.abi3.so": "mhlZCzeFfCZNZ0/XaMGpj9x0zU3vYjT314g1dbysmxM=",
+      "_ops.py": "XOwQje24M0e2Is75hhdYmSaw60wjwb/z4WY2Fza5UIc="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "81f55ea30fd8f819dcf93a3c934dd584c895bd2f",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "85897539ddf848f1cbd4bb17d24c89030d2c33e2",
+      "dirty": false
+    }
   }
+}

build/torch213-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,305 @@

+"""FlashRT causal Conv1D state kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _empty_bias_like(x: torch.Tensor) -> torch.Tensor:
+    return torch.empty((0,), device=x.device, dtype=torch.bfloat16)
+def _check_conv_shapes(x: torch.Tensor, w: torch.Tensor, bias: torch.Tensor, has_bias: bool, out: torch.Tensor) -> None:
+    if x.dim() != 3:
+        raise RuntimeError("x must have shape (B,S,C)")
+    if w.dim() != 2 or w.shape[0] != x.shape[2] or w.shape[1] < 2 or w.shape[1] > 8:
+        raise RuntimeError("w must have shape (C,K), 2 <= K <= 8")
+    if has_bias and bias.shape != (x.shape[2],):
+        raise RuntimeError("bias must have shape (C,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must match x shape")
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_bf16"))
+def _causal_conv1d_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    _check_conv_shapes(x, w, bias, has_bias, out)
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_bf16"))
+def _causal_conv1d_update_bf16_fake(
+    x_new: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    state: torch.Tensor,
+    out: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    if x_new.dim() != 2:
+        raise RuntimeError("x_new must have shape (B,C)")
+    b, c = x_new.shape
+    if w.dim() != 2 or w.shape[0] != c or w.shape[1] < 2 or w.shape[1] > 8:
+        raise RuntimeError("w must have shape (C,K), 2 <= K <= 8")
+    if has_bias and bias.shape != (c,):
+        raise RuntimeError("bias must have shape (C,)")
+    if state.shape != (b, c, w.shape[1] - 1):
+        raise RuntimeError("state must have shape (B,C,K-1)")
+    if out.shape != x_new.shape:
+        raise RuntimeError("out must match x_new shape")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_inout_bf16"))
+def _causal_conv1d_update_inout_bf16_fake(
+    x_new: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    state_in: torch.Tensor,
+    state_out: torch.Tensor,
+    out: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    _causal_conv1d_update_bf16_fake(x_new, w, bias, state_in, out, has_bias, apply_silu)
+    if state_out.shape != state_in.shape:
+        raise RuntimeError("state_out shape mismatch")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_chunk_bf16"))
+def _causal_conv1d_update_chunk_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    state: torch.Tensor,
+    out: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    _check_conv_shapes(x, w, bias, has_bias, out)
+    if state.shape != (x.shape[0], x.shape[2], w.shape[1] - 1):
+        raise RuntimeError("state must have shape (B,C,K-1)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_chunk_parallel_bf16"))
+def _causal_conv1d_update_chunk_parallel_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    state: torch.Tensor,
+    out: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    _causal_conv1d_update_chunk_bf16_fake(x, w, bias, state, out, has_bias, apply_silu)
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_chunk_parallel_gqa_bf16"))
+def _causal_conv1d_update_chunk_parallel_gqa_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    state: torch.Tensor,
+    q16: torch.Tensor,
+    k16: torch.Tensor,
+    v48: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    if x.dim() != 3 or x.shape[2] != 10240:
+        raise RuntimeError("GQA split variant expects x shape (B,S,10240)")
+    if w.shape != (10240, w.shape[1]) or w.shape[1] < 2 or w.shape[1] > 8:
+        raise RuntimeError("w must have shape (10240,K), 2 <= K <= 8")
+    b, s, _ = x.shape
+    if state.shape != (b, 10240, w.shape[1] - 1):
+        raise RuntimeError("state must have shape (B,10240,K-1)")
+    if q16.shape != (b, s, 16, 128) or k16.shape != q16.shape or v48.shape != (b, s, 48, 128):
+        raise RuntimeError("q16/k16/v48 output shape mismatch")
+    if has_bias and bias.shape != (10240,):
+        raise RuntimeError("bias must have shape (10240,)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_steps_gqa_bf16"))
+def _causal_conv1d_update_steps_gqa_bf16_fake(
+    x, w, bias, state, q16, k16, v48, apply_silu: bool = True
+) -> None:
+    del apply_silu
+    if x.dim() != 2 or x.shape[1] != 10240:
+        raise RuntimeError("x must have shape (S,10240)")
+    s = x.shape[0]
+    if w.shape != (10240, 4) or bias.shape != (10240,) or state.shape != (10240, 3):
+        raise RuntimeError("w/bias/state shape contract failed")
+    if q16.shape != (s, 2048) or k16.shape != q16.shape or v48.shape != (s, 6144):
+        raise RuntimeError("q16/k16/v48 output shape mismatch")
+    return None
+def causal_conv1d_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(x)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x)
+    ops.causal_conv1d_bf16(x, w, bias, out, bool(has_bias), bool(apply_silu))
+    return out
+def causal_conv1d_update_bf16(
+    x_new: torch.Tensor,
+    w: torch.Tensor,
+    state: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(x_new)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x_new)
+    ops.causal_conv1d_update_bf16(x_new, w, bias, state, out, bool(has_bias), bool(apply_silu))
+    return out
+def causal_conv1d_update_inout_bf16(
+    x_new: torch.Tensor,
+    w: torch.Tensor,
+    state_in: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    out: Optional[torch.Tensor] = None,
+    state_out: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if out is None:
+        out = torch.empty_like(x_new)
+    if state_out is None:
+        state_out = torch.empty_like(state_in)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x_new)
+    ops.causal_conv1d_update_inout_bf16(
+        x_new, w, bias, state_in, state_out, out, bool(has_bias), bool(apply_silu)
+    )
+    return out, state_out
+def causal_conv1d_update_chunk_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    state: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(x)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x)
+    ops.causal_conv1d_update_chunk_bf16(x, w, bias, state, out, bool(has_bias), bool(apply_silu))
+    return out
+def causal_conv1d_update_chunk_parallel_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    state: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(x)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x)
+    ops.causal_conv1d_update_chunk_parallel_bf16(x, w, bias, state, out, bool(has_bias), bool(apply_silu))
+    return out
+def causal_conv1d_update_chunk_parallel_gqa_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    state: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    q16: Optional[torch.Tensor] = None,
+    k16: Optional[torch.Tensor] = None,
+    v48: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    b, s, _ = x.shape
+    if q16 is None:
+        q16 = torch.empty((b, s, 16, 128), device=x.device, dtype=torch.bfloat16)
+    if k16 is None:
+        k16 = torch.empty_like(q16)
+    if v48 is None:
+        v48 = torch.empty((b, s, 48, 128), device=x.device, dtype=torch.bfloat16)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x)
+    ops.causal_conv1d_update_chunk_parallel_gqa_bf16(
+        x, w, bias, state, q16, k16, v48, bool(has_bias), bool(apply_silu)
+    )
+    return q16, k16, v48
+def causal_conv1d_update_steps_gqa_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    state: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    apply_silu: bool = True,
+    q16: Optional[torch.Tensor] = None,
+    k16: Optional[torch.Tensor] = None,
+    v48: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    s = x.shape[0]
+    if q16 is None:
+        q16 = torch.empty((s, 2048), device=x.device, dtype=torch.bfloat16)
+    if k16 is None:
+        k16 = torch.empty_like(q16)
+    if v48 is None:
+        v48 = torch.empty((s, 6144), device=x.device, dtype=torch.bfloat16)
+    ops.causal_conv1d_update_steps_gqa_bf16(
+        x, w, bias, state, q16, k16, v48, bool(apply_silu)
+    )
+    return q16, k16, v48
+__all__ = [
+    "causal_conv1d_bf16",
+    "causal_conv1d_update_bf16",
+    "causal_conv1d_update_inout_bf16",
+    "causal_conv1d_update_chunk_bf16",
+    "causal_conv1d_update_chunk_parallel_bf16",
+    "causal_conv1d_update_chunk_parallel_gqa_bf16",
+    "causal_conv1d_update_steps_gqa_bf16",
+]

build/torch213-cxx11-cu130-x86_64-linux/_causal_conv1d_state_cuda_8589753.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:449890100502e5de595e747fafb3c9b2816255b8780754cf9aa362b244a67d25
+size 1256232

build/torch213-cxx11-cu130-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _causal_conv1d_state_cuda_8589753
+ops = torch.ops._causal_conv1d_state_cuda_8589753
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_causal_conv1d_state_cuda_8589753::{op_name}"

build/torch213-cxx11-cu130-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,41 @@

+{
+  "name": "causal-conv1d-state",
+  "id": "_causal_conv1d_state_cuda_8589753",
+  "version": 2,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "jC3cFEzRQxoloALJMANJU4Px3BpryCaxyVp+MdmTj6Q=",
+      "_causal_conv1d_state_cuda_8589753.abi3.so": "RJiQEAUC5d5ZXnR/r7PJsoFiVbh4B1TPmqNiskSmfSU=",
+      "_ops.py": "XOwQje24M0e2Is75hhdYmSaw60wjwb/z4WY2Fza5UIc="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "81f55ea30fd8f819dcf93a3c934dd584c895bd2f",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "85897539ddf848f1cbd4bb17d24c89030d2c33e2",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu132-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,305 @@

+"""FlashRT causal Conv1D state kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _empty_bias_like(x: torch.Tensor) -> torch.Tensor:
+    return torch.empty((0,), device=x.device, dtype=torch.bfloat16)
+def _check_conv_shapes(x: torch.Tensor, w: torch.Tensor, bias: torch.Tensor, has_bias: bool, out: torch.Tensor) -> None:
+    if x.dim() != 3:
+        raise RuntimeError("x must have shape (B,S,C)")
+    if w.dim() != 2 or w.shape[0] != x.shape[2] or w.shape[1] < 2 or w.shape[1] > 8:
+        raise RuntimeError("w must have shape (C,K), 2 <= K <= 8")
+    if has_bias and bias.shape != (x.shape[2],):
+        raise RuntimeError("bias must have shape (C,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must match x shape")
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_bf16"))
+def _causal_conv1d_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    _check_conv_shapes(x, w, bias, has_bias, out)
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_bf16"))
+def _causal_conv1d_update_bf16_fake(
+    x_new: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    state: torch.Tensor,
+    out: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    if x_new.dim() != 2:
+        raise RuntimeError("x_new must have shape (B,C)")
+    b, c = x_new.shape
+    if w.dim() != 2 or w.shape[0] != c or w.shape[1] < 2 or w.shape[1] > 8:
+        raise RuntimeError("w must have shape (C,K), 2 <= K <= 8")
+    if has_bias and bias.shape != (c,):
+        raise RuntimeError("bias must have shape (C,)")
+    if state.shape != (b, c, w.shape[1] - 1):
+        raise RuntimeError("state must have shape (B,C,K-1)")
+    if out.shape != x_new.shape:
+        raise RuntimeError("out must match x_new shape")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_inout_bf16"))
+def _causal_conv1d_update_inout_bf16_fake(
+    x_new: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    state_in: torch.Tensor,
+    state_out: torch.Tensor,
+    out: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    _causal_conv1d_update_bf16_fake(x_new, w, bias, state_in, out, has_bias, apply_silu)
+    if state_out.shape != state_in.shape:
+        raise RuntimeError("state_out shape mismatch")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_chunk_bf16"))
+def _causal_conv1d_update_chunk_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    state: torch.Tensor,
+    out: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    _check_conv_shapes(x, w, bias, has_bias, out)
+    if state.shape != (x.shape[0], x.shape[2], w.shape[1] - 1):
+        raise RuntimeError("state must have shape (B,C,K-1)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_chunk_parallel_bf16"))
+def _causal_conv1d_update_chunk_parallel_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    state: torch.Tensor,
+    out: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    _causal_conv1d_update_chunk_bf16_fake(x, w, bias, state, out, has_bias, apply_silu)
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_chunk_parallel_gqa_bf16"))
+def _causal_conv1d_update_chunk_parallel_gqa_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    state: torch.Tensor,
+    q16: torch.Tensor,
+    k16: torch.Tensor,
+    v48: torch.Tensor,
+    has_bias: bool = True,
+    apply_silu: bool = True,
+) -> None:
+    if x.dim() != 3 or x.shape[2] != 10240:
+        raise RuntimeError("GQA split variant expects x shape (B,S,10240)")
+    if w.shape != (10240, w.shape[1]) or w.shape[1] < 2 or w.shape[1] > 8:
+        raise RuntimeError("w must have shape (10240,K), 2 <= K <= 8")
+    b, s, _ = x.shape
+    if state.shape != (b, 10240, w.shape[1] - 1):
+        raise RuntimeError("state must have shape (B,10240,K-1)")
+    if q16.shape != (b, s, 16, 128) or k16.shape != q16.shape or v48.shape != (b, s, 48, 128):
+        raise RuntimeError("q16/k16/v48 output shape mismatch")
+    if has_bias and bias.shape != (10240,):
+        raise RuntimeError("bias must have shape (10240,)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("causal_conv1d_update_steps_gqa_bf16"))
+def _causal_conv1d_update_steps_gqa_bf16_fake(
+    x, w, bias, state, q16, k16, v48, apply_silu: bool = True
+) -> None:
+    del apply_silu
+    if x.dim() != 2 or x.shape[1] != 10240:
+        raise RuntimeError("x must have shape (S,10240)")
+    s = x.shape[0]
+    if w.shape != (10240, 4) or bias.shape != (10240,) or state.shape != (10240, 3):
+        raise RuntimeError("w/bias/state shape contract failed")
+    if q16.shape != (s, 2048) or k16.shape != q16.shape or v48.shape != (s, 6144):
+        raise RuntimeError("q16/k16/v48 output shape mismatch")
+    return None
+def causal_conv1d_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(x)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x)
+    ops.causal_conv1d_bf16(x, w, bias, out, bool(has_bias), bool(apply_silu))
+    return out
+def causal_conv1d_update_bf16(
+    x_new: torch.Tensor,
+    w: torch.Tensor,
+    state: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(x_new)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x_new)
+    ops.causal_conv1d_update_bf16(x_new, w, bias, state, out, bool(has_bias), bool(apply_silu))
+    return out
+def causal_conv1d_update_inout_bf16(
+    x_new: torch.Tensor,
+    w: torch.Tensor,
+    state_in: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    out: Optional[torch.Tensor] = None,
+    state_out: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if out is None:
+        out = torch.empty_like(x_new)
+    if state_out is None:
+        state_out = torch.empty_like(state_in)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x_new)
+    ops.causal_conv1d_update_inout_bf16(
+        x_new, w, bias, state_in, state_out, out, bool(has_bias), bool(apply_silu)
+    )
+    return out, state_out
+def causal_conv1d_update_chunk_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    state: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(x)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x)
+    ops.causal_conv1d_update_chunk_bf16(x, w, bias, state, out, bool(has_bias), bool(apply_silu))
+    return out
+def causal_conv1d_update_chunk_parallel_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    state: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(x)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x)
+    ops.causal_conv1d_update_chunk_parallel_bf16(x, w, bias, state, out, bool(has_bias), bool(apply_silu))
+    return out
+def causal_conv1d_update_chunk_parallel_gqa_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    state: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    *,
+    apply_silu: bool = True,
+    q16: Optional[torch.Tensor] = None,
+    k16: Optional[torch.Tensor] = None,
+    v48: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    b, s, _ = x.shape
+    if q16 is None:
+        q16 = torch.empty((b, s, 16, 128), device=x.device, dtype=torch.bfloat16)
+    if k16 is None:
+        k16 = torch.empty_like(q16)
+    if v48 is None:
+        v48 = torch.empty((b, s, 48, 128), device=x.device, dtype=torch.bfloat16)
+    has_bias = bias is not None
+    if bias is None:
+        bias = _empty_bias_like(x)
+    ops.causal_conv1d_update_chunk_parallel_gqa_bf16(
+        x, w, bias, state, q16, k16, v48, bool(has_bias), bool(apply_silu)
+    )
+    return q16, k16, v48
+def causal_conv1d_update_steps_gqa_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    state: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    apply_silu: bool = True,
+    q16: Optional[torch.Tensor] = None,
+    k16: Optional[torch.Tensor] = None,
+    v48: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    s = x.shape[0]
+    if q16 is None:
+        q16 = torch.empty((s, 2048), device=x.device, dtype=torch.bfloat16)
+    if k16 is None:
+        k16 = torch.empty_like(q16)
+    if v48 is None:
+        v48 = torch.empty((s, 6144), device=x.device, dtype=torch.bfloat16)
+    ops.causal_conv1d_update_steps_gqa_bf16(
+        x, w, bias, state, q16, k16, v48, bool(apply_silu)
+    )
+    return q16, k16, v48
+__all__ = [
+    "causal_conv1d_bf16",
+    "causal_conv1d_update_bf16",
+    "causal_conv1d_update_inout_bf16",
+    "causal_conv1d_update_chunk_bf16",
+    "causal_conv1d_update_chunk_parallel_bf16",
+    "causal_conv1d_update_chunk_parallel_gqa_bf16",
+    "causal_conv1d_update_steps_gqa_bf16",
+]

build/torch213-cxx11-cu132-x86_64-linux/_causal_conv1d_state_cuda_8589753.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:fefab8142cd9d244dd948c26905edc8d689fceeb70b0ff9a2423c3fea23b5c2e
+size 1264424

build/torch213-cxx11-cu132-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _causal_conv1d_state_cuda_8589753
+ops = torch.ops._causal_conv1d_state_cuda_8589753
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_causal_conv1d_state_cuda_8589753::{op_name}"

build/torch213-cxx11-cu132-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,41 @@

+{
+  "name": "causal-conv1d-state",
+  "id": "_causal_conv1d_state_cuda_8589753",
+  "version": 2,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "jC3cFEzRQxoloALJMANJU4Px3BpryCaxyVp+MdmTj6Q=",
+      "_causal_conv1d_state_cuda_8589753.abi3.so": "/vq4FCzZ0kTdlIwmkF7cjWifzutwsP+aJCPD/qI7XC4=",
+      "_ops.py": "XOwQje24M0e2Is75hhdYmSaw60wjwb/z4WY2Fza5UIc="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "81f55ea30fd8f819dcf93a3c934dd584c895bd2f",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "85897539ddf848f1cbd4bb17d24c89030d2c33e2",
+      "dirty": false
+    }
+  }
+}