Uploaded using `kernel-builder`.

Browse files

Files changed (15) hide show

build/torch211-cxx11-cpu-x86_64-linux/__init__.py +90 -0
build/torch211-cxx11-cpu-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cpu-x86_64-linux/_spec_verify_cpu_9f08509.abi3.so +3 -0
build/torch211-cxx11-cpu-x86_64-linux/metadata.json +30 -0
build/torch211-cxx11-cpu-x86_64-linux/spec_verify/__init__.py +26 -0
build/torch212-cxx11-cpu-x86_64-linux/__init__.py +90 -0
build/torch212-cxx11-cpu-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cpu-x86_64-linux/_spec_verify_cpu_9f08509.abi3.so +3 -0
build/torch212-cxx11-cpu-x86_64-linux/metadata.json +30 -0
build/torch212-cxx11-cpu-x86_64-linux/spec_verify/__init__.py +26 -0
build/torch213-cxx11-cpu-x86_64-linux/__init__.py +90 -0
build/torch213-cxx11-cpu-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cpu-x86_64-linux/_spec_verify_cpu_9f08509.abi3.so +3 -0
build/torch213-cxx11-cpu-x86_64-linux/metadata.json +30 -0
build/torch213-cxx11-cpu-x86_64-linux/spec_verify/__init__.py +26 -0

build/torch211-cxx11-cpu-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,90 @@

+"""spec-verify: speculative-decoding verification ops for CPUs, aarch64-first.
+A draft model proposes K tokens; the target verifies them in one [K+1, V]
+forward (the batched-GEMM path where the companion linear kernels run
+fastest). These ops turn that forward into an accepted prefix plus one
+correction token: greedy (temperature 0) verification, or distribution-
+preserving speculative sampling. Companions on the Hub: bitnet-cpu,
+quant-matmul, gguf-gemv (linears), cpu-attn, decode-ops.
+High-level API:
+    verify_greedy(target_probs, draft_tokens) -> (n_accepted, next_token)
+    accept(target_probs, draft_probs, draft_tokens, generator)
+    SpeculativeDecoder(draft_step, target_step, ...)   orchestration loop
+"""
+from typing import Callable, List, Optional, Tuple
+import torch
+from ._ops import ops
+def verify_greedy(target_probs: torch.Tensor,
+                  draft_tokens: torch.Tensor) -> Tuple[int, int]:
+    """Temperature-0 verify. target_probs [K+1, V] (probs or logits; only the
+    argmax is used), draft_tokens [K] int64. Returns (n_accepted, next_token):
+    keep draft_tokens[:n_accepted] and append next_token."""
+    tp = target_probs.to(torch.float32).contiguous()
+    dt = draft_tokens.to(torch.int64).contiguous()
+    n, tok = ops.sv_verify_greedy(tp, dt)
+    return int(n), int(tok)
+def accept(target_probs: torch.Tensor, draft_probs: torch.Tensor,
+           draft_tokens: torch.Tensor,
+           generator: Optional[torch.Generator] = None) -> Tuple[int, int]:
+    """Distribution-preserving speculative sampling.
+    target_probs [K+1, V] and draft_probs [K, V] must be post-warp
+    probability distributions warped identically (the caller's contract).
+    Returns (n_accepted, next_token), deterministic given the generator.
+    """
+    K = int(draft_tokens.numel())
+    ra = torch.rand(K, generator=generator, dtype=torch.float32)
+    rr = float(torch.rand(1, generator=generator).item())
+    tp = target_probs.to(torch.float32).contiguous()
+    dp = draft_probs.to(torch.float32).contiguous()
+    dt = draft_tokens.to(torch.int64).contiguous()
+    n, tok = ops.sv_accept(tp, dp, dt, ra.contiguous(), rr)
+    return int(n), int(tok)
+class SpeculativeDecoder:
+    """Draft/verify loop over user-supplied step callbacks.
+    draft_step(prefix, k) -> (tokens [k] int64, probs [k, V] or None)
+    target_step(prefix, proposed [k]) -> target_probs [k+1, V]
+    With temperature 0 (greedy) draft probs are unused; otherwise the caller
+    returns post-warp draft probabilities and the target_probs must carry the
+    same warping for the sampling to preserve the target distribution.
+    """
+    def __init__(self, draft_step: Callable, target_step: Callable,
+                 k: int = 4, greedy: bool = True,
+                 generator: Optional[torch.Generator] = None):
+        self.draft_step = draft_step
+        self.target_step = target_step
+        self.k = k
+        self.greedy = greedy
+        self.generator = generator
+    def step(self, prefix: List[int]) -> Tuple[List[int], int]:
+        """One speculative round: returns (accepted_tokens, n_target_calls=1)."""
+        proposed, dprobs = self.draft_step(prefix, self.k)
+        tprobs = self.target_step(prefix, proposed)
+        if self.greedy:
+            n, tok = verify_greedy(tprobs, proposed)
+        else:
+            n, tok = accept(tprobs, dprobs, proposed, self.generator)
+        return [int(proposed[i]) for i in range(n)] + [tok], 1
+    def generate(self, prefix: List[int], max_new: int) -> List[int]:
+        out = list(prefix)
+        while len(out) - len(prefix) < max_new:
+            accepted, _ = self.step(out)
+            out.extend(accepted)
+        return out[: len(prefix) + max_new]
+__all__ = ["verify_greedy", "accept", "SpeculativeDecoder"]

build/torch211-cxx11-cpu-x86_64-linux/_ops.py ADDED Viewed

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

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

build/torch211-cxx11-cpu-x86_64-linux/_spec_verify_cpu_9f08509.abi3.so ADDED Viewed

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

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

build/torch211-cxx11-cpu-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,30 @@

+{
+  "name": "spec-verify",
+  "id": "_spec_verify_cpu_9f08509",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cpu"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "UlLRZztEids73FkPtfMh14gPWORZUqSZ6p1k/bas1pI=",
+      "_ops.py": "Cz9lzcNnPDNcX6cbRzRg5c5NfXKMaGHpqCA0XFdtthE=",
+      "_spec_verify_cpu_9f08509.abi3.so": "nm62ocBz6ogjMgAZrVupYGgPFC6p+lga/St808OZEqI=",
+      "spec_verify/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "9f085096c7c6f9326be6bdb489d64f8ae81278a9",
+      "dirty": false
+    }
+  }
+}

build/torch211-cxx11-cpu-x86_64-linux/spec_verify/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+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-cpu-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,90 @@

+"""spec-verify: speculative-decoding verification ops for CPUs, aarch64-first.
+A draft model proposes K tokens; the target verifies them in one [K+1, V]
+forward (the batched-GEMM path where the companion linear kernels run
+fastest). These ops turn that forward into an accepted prefix plus one
+correction token: greedy (temperature 0) verification, or distribution-
+preserving speculative sampling. Companions on the Hub: bitnet-cpu,
+quant-matmul, gguf-gemv (linears), cpu-attn, decode-ops.
+High-level API:
+    verify_greedy(target_probs, draft_tokens) -> (n_accepted, next_token)
+    accept(target_probs, draft_probs, draft_tokens, generator)
+    SpeculativeDecoder(draft_step, target_step, ...)   orchestration loop
+"""
+from typing import Callable, List, Optional, Tuple
+import torch
+from ._ops import ops
+def verify_greedy(target_probs: torch.Tensor,
+                  draft_tokens: torch.Tensor) -> Tuple[int, int]:
+    """Temperature-0 verify. target_probs [K+1, V] (probs or logits; only the
+    argmax is used), draft_tokens [K] int64. Returns (n_accepted, next_token):
+    keep draft_tokens[:n_accepted] and append next_token."""
+    tp = target_probs.to(torch.float32).contiguous()
+    dt = draft_tokens.to(torch.int64).contiguous()
+    n, tok = ops.sv_verify_greedy(tp, dt)
+    return int(n), int(tok)
+def accept(target_probs: torch.Tensor, draft_probs: torch.Tensor,
+           draft_tokens: torch.Tensor,
+           generator: Optional[torch.Generator] = None) -> Tuple[int, int]:
+    """Distribution-preserving speculative sampling.
+    target_probs [K+1, V] and draft_probs [K, V] must be post-warp
+    probability distributions warped identically (the caller's contract).
+    Returns (n_accepted, next_token), deterministic given the generator.
+    """
+    K = int(draft_tokens.numel())
+    ra = torch.rand(K, generator=generator, dtype=torch.float32)
+    rr = float(torch.rand(1, generator=generator).item())
+    tp = target_probs.to(torch.float32).contiguous()
+    dp = draft_probs.to(torch.float32).contiguous()
+    dt = draft_tokens.to(torch.int64).contiguous()
+    n, tok = ops.sv_accept(tp, dp, dt, ra.contiguous(), rr)
+    return int(n), int(tok)
+class SpeculativeDecoder:
+    """Draft/verify loop over user-supplied step callbacks.
+    draft_step(prefix, k) -> (tokens [k] int64, probs [k, V] or None)
+    target_step(prefix, proposed [k]) -> target_probs [k+1, V]
+    With temperature 0 (greedy) draft probs are unused; otherwise the caller
+    returns post-warp draft probabilities and the target_probs must carry the
+    same warping for the sampling to preserve the target distribution.
+    """
+    def __init__(self, draft_step: Callable, target_step: Callable,
+                 k: int = 4, greedy: bool = True,
+                 generator: Optional[torch.Generator] = None):
+        self.draft_step = draft_step
+        self.target_step = target_step
+        self.k = k
+        self.greedy = greedy
+        self.generator = generator
+    def step(self, prefix: List[int]) -> Tuple[List[int], int]:
+        """One speculative round: returns (accepted_tokens, n_target_calls=1)."""
+        proposed, dprobs = self.draft_step(prefix, self.k)
+        tprobs = self.target_step(prefix, proposed)
+        if self.greedy:
+            n, tok = verify_greedy(tprobs, proposed)
+        else:
+            n, tok = accept(tprobs, dprobs, proposed, self.generator)
+        return [int(proposed[i]) for i in range(n)] + [tok], 1
+    def generate(self, prefix: List[int], max_new: int) -> List[int]:
+        out = list(prefix)
+        while len(out) - len(prefix) < max_new:
+            accepted, _ = self.step(out)
+            out.extend(accepted)
+        return out[: len(prefix) + max_new]
+__all__ = ["verify_greedy", "accept", "SpeculativeDecoder"]

build/torch212-cxx11-cpu-x86_64-linux/_ops.py ADDED Viewed

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

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

build/torch212-cxx11-cpu-x86_64-linux/_spec_verify_cpu_9f08509.abi3.so ADDED Viewed

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

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

build/torch212-cxx11-cpu-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,30 @@

+{
+  "name": "spec-verify",
+  "id": "_spec_verify_cpu_9f08509",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cpu"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "UlLRZztEids73FkPtfMh14gPWORZUqSZ6p1k/bas1pI=",
+      "_ops.py": "Cz9lzcNnPDNcX6cbRzRg5c5NfXKMaGHpqCA0XFdtthE=",
+      "_spec_verify_cpu_9f08509.abi3.so": "xxc2+A2AoFnLs8+F7dhEUA5Qzr3NU6GRsFQFcBlpYis=",
+      "spec_verify/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "9f085096c7c6f9326be6bdb489d64f8ae81278a9",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cpu-x86_64-linux/spec_verify/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+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/torch213-cxx11-cpu-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,90 @@

+"""spec-verify: speculative-decoding verification ops for CPUs, aarch64-first.
+A draft model proposes K tokens; the target verifies them in one [K+1, V]
+forward (the batched-GEMM path where the companion linear kernels run
+fastest). These ops turn that forward into an accepted prefix plus one
+correction token: greedy (temperature 0) verification, or distribution-
+preserving speculative sampling. Companions on the Hub: bitnet-cpu,
+quant-matmul, gguf-gemv (linears), cpu-attn, decode-ops.
+High-level API:
+    verify_greedy(target_probs, draft_tokens) -> (n_accepted, next_token)
+    accept(target_probs, draft_probs, draft_tokens, generator)
+    SpeculativeDecoder(draft_step, target_step, ...)   orchestration loop
+"""
+from typing import Callable, List, Optional, Tuple
+import torch
+from ._ops import ops
+def verify_greedy(target_probs: torch.Tensor,
+                  draft_tokens: torch.Tensor) -> Tuple[int, int]:
+    """Temperature-0 verify. target_probs [K+1, V] (probs or logits; only the
+    argmax is used), draft_tokens [K] int64. Returns (n_accepted, next_token):
+    keep draft_tokens[:n_accepted] and append next_token."""
+    tp = target_probs.to(torch.float32).contiguous()
+    dt = draft_tokens.to(torch.int64).contiguous()
+    n, tok = ops.sv_verify_greedy(tp, dt)
+    return int(n), int(tok)
+def accept(target_probs: torch.Tensor, draft_probs: torch.Tensor,
+           draft_tokens: torch.Tensor,
+           generator: Optional[torch.Generator] = None) -> Tuple[int, int]:
+    """Distribution-preserving speculative sampling.
+    target_probs [K+1, V] and draft_probs [K, V] must be post-warp
+    probability distributions warped identically (the caller's contract).
+    Returns (n_accepted, next_token), deterministic given the generator.
+    """
+    K = int(draft_tokens.numel())
+    ra = torch.rand(K, generator=generator, dtype=torch.float32)
+    rr = float(torch.rand(1, generator=generator).item())
+    tp = target_probs.to(torch.float32).contiguous()
+    dp = draft_probs.to(torch.float32).contiguous()
+    dt = draft_tokens.to(torch.int64).contiguous()
+    n, tok = ops.sv_accept(tp, dp, dt, ra.contiguous(), rr)
+    return int(n), int(tok)
+class SpeculativeDecoder:
+    """Draft/verify loop over user-supplied step callbacks.
+    draft_step(prefix, k) -> (tokens [k] int64, probs [k, V] or None)
+    target_step(prefix, proposed [k]) -> target_probs [k+1, V]
+    With temperature 0 (greedy) draft probs are unused; otherwise the caller
+    returns post-warp draft probabilities and the target_probs must carry the
+    same warping for the sampling to preserve the target distribution.
+    """
+    def __init__(self, draft_step: Callable, target_step: Callable,
+                 k: int = 4, greedy: bool = True,
+                 generator: Optional[torch.Generator] = None):
+        self.draft_step = draft_step
+        self.target_step = target_step
+        self.k = k
+        self.greedy = greedy
+        self.generator = generator
+    def step(self, prefix: List[int]) -> Tuple[List[int], int]:
+        """One speculative round: returns (accepted_tokens, n_target_calls=1)."""
+        proposed, dprobs = self.draft_step(prefix, self.k)
+        tprobs = self.target_step(prefix, proposed)
+        if self.greedy:
+            n, tok = verify_greedy(tprobs, proposed)
+        else:
+            n, tok = accept(tprobs, dprobs, proposed, self.generator)
+        return [int(proposed[i]) for i in range(n)] + [tok], 1
+    def generate(self, prefix: List[int], max_new: int) -> List[int]:
+        out = list(prefix)
+        while len(out) - len(prefix) < max_new:
+            accepted, _ = self.step(out)
+            out.extend(accepted)
+        return out[: len(prefix) + max_new]
+__all__ = ["verify_greedy", "accept", "SpeculativeDecoder"]

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

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

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

build/torch213-cxx11-cpu-x86_64-linux/_spec_verify_cpu_9f08509.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:0c1a24ef1dd8b94ee1901d0b42d2ff727193ac535cac505fc5fbe5529802bf41
+size 70096

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

	@@ -0,0 +1,30 @@

+{
+  "name": "spec-verify",
+  "id": "_spec_verify_cpu_9f08509",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cpu"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "UlLRZztEids73FkPtfMh14gPWORZUqSZ6p1k/bas1pI=",
+      "_ops.py": "Cz9lzcNnPDNcX6cbRzRg5c5NfXKMaGHpqCA0XFdtthE=",
+      "_spec_verify_cpu_9f08509.abi3.so": "DBok7x3YuU7hkB0LQtL/cnGTrFNcrFBfxfvlUpgCv0E=",
+      "spec_verify/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "9f085096c7c6f9326be6bdb489d64f8ae81278a9",
+      "dirty": false
+    }
+  }
+}

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

	@@ -0,0 +1,26 @@

+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")))