Add torch211-cxx11-cu130-aarch64-linux SM110 artifact

Browse files

Files changed (5) hide show

build/torch211-cxx11-cu130-aarch64-linux/__init__.py +105 -0
build/torch211-cxx11-cu130-aarch64-linux/_ops.py +6 -0
build/torch211-cxx11-cu130-aarch64-linux/_speculative_draft_primitives_cuda_6ee7cab.abi3.so +3 -0
build/torch211-cxx11-cu130-aarch64-linux/metadata.json +22 -0
build/torch211-cxx11-cu130-aarch64-linux/speculative_draft_primitives/__init__.py +14 -0

build/torch211-cxx11-cu130-aarch64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,105 @@

+"""FlashRT speculative decoding helper kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(add_op_namespace_prefix("argmax_bf16"))
+def _argmax_bf16_fake(logits: torch.Tensor, argmax_out: torch.Tensor) -> None:
+    if logits.dim() != 2 or argmax_out.shape != (logits.shape[0],):
+        raise RuntimeError("argmax_bf16 expects logits (rows,vocab), argmax_out (rows,)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("accept_greedy_bf16"))
+def _accept_greedy_bf16_fake(
+    logits: torch.Tensor,
+    drafts: torch.Tensor,
+    argmax_out: torch.Tensor,
+    accept_n: torch.Tensor,
+    spec_k: int,
+) -> None:
+    if logits.dim() != 2 or argmax_out.shape != (logits.shape[0],):
+        raise RuntimeError("accept_greedy_bf16 expects logits (rows,vocab), argmax_out (rows,)")
+    if drafts.dim() != 1 or drafts.numel() < spec_k or accept_n.numel() < 1:
+        raise RuntimeError("drafts/accept_n shape mismatch")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("accept_partitioned_bf16"))
+def _accept_partitioned_bf16_fake(
+    logits: torch.Tensor,
+    drafts: torch.Tensor,
+    argmax_out: torch.Tensor,
+    accept_n: torch.Tensor,
+    partial_vals: torch.Tensor,
+    partial_idx: torch.Tensor,
+    spec_k: int,
+    parts: int,
+) -> None:
+    if partial_vals.shape != (logits.shape[0], parts) or partial_idx.shape != (logits.shape[0], parts):
+        raise RuntimeError("partial buffers must have shape (rows, parts)")
+    return _accept_greedy_bf16_fake(logits, drafts, argmax_out, accept_n, spec_k)
+def argmax_bf16(logits: torch.Tensor, *, out: Optional[torch.Tensor] = None) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((logits.shape[0],), device=logits.device, dtype=torch.int64)
+    ops.argmax_bf16(logits, out)
+    return out
+def accept_greedy_bf16(
+    logits: torch.Tensor,
+    drafts: torch.Tensor,
+    spec_k: int,
+    *,
+    argmax_out: Optional[torch.Tensor] = None,
+    accept_n: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if argmax_out is None:
+        argmax_out = torch.empty((logits.shape[0],), device=logits.device, dtype=torch.int64)
+    if accept_n is None:
+        accept_n = torch.empty((1,), device=logits.device, dtype=torch.int32)
+    ops.accept_greedy_bf16(logits, drafts, argmax_out, accept_n, int(spec_k))
+    return argmax_out, accept_n
+def accept_partitioned_bf16(
+    logits: torch.Tensor,
+    drafts: torch.Tensor,
+    spec_k: int,
+    parts: Optional[int] = None,
+    *,
+    argmax_out: Optional[torch.Tensor] = None,
+    accept_n: Optional[torch.Tensor] = None,
+    partial_vals: Optional[torch.Tensor] = None,
+    partial_idx: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if parts is None:
+        vocab = int(logits.shape[1])
+        parts = 32 if vocab >= 131072 else (16 if vocab >= 65536 else 8)
+    if argmax_out is None:
+        argmax_out = torch.empty((logits.shape[0],), device=logits.device, dtype=torch.int64)
+    if accept_n is None:
+        accept_n = torch.empty((1,), device=logits.device, dtype=torch.int32)
+    if partial_vals is None:
+        partial_vals = torch.empty((logits.shape[0], parts), device=logits.device, dtype=torch.float32)
+    if partial_idx is None:
+        partial_idx = torch.empty((logits.shape[0], parts), device=logits.device, dtype=torch.int32)
+    ops.accept_partitioned_bf16(
+        logits, drafts, argmax_out, accept_n, partial_vals, partial_idx, int(spec_k), int(parts)
+    )
+    return argmax_out, accept_n
+__all__ = [
+    "argmax_bf16",
+    "accept_greedy_bf16",
+    "accept_partitioned_bf16",
+]

build/torch211-cxx11-cu130-aarch64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,6 @@

+import torch
+from . import _speculative_draft_primitives_cuda_6ee7cab
+ops = torch.ops._speculative_draft_primitives_cuda_6ee7cab
+def add_op_namespace_prefix(op_name: str):
+    return f"_speculative_draft_primitives_cuda_6ee7cab::{op_name}"

build/torch211-cxx11-cu130-aarch64-linux/_speculative_draft_primitives_cuda_6ee7cab.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:6b95428d63268aaf762433b0bbf183219fa0571811edb8c31ea483233f41a60a
+size 175736

build/torch211-cxx11-cu130-aarch64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,22 @@

+{
+  "name": "speculative-draft-primitives",
+  "id": "_speculative_draft_primitives_cuda_6ee7cab",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "6NhsYcUvz+6W9eEWWCQJVI8aaXA89WhNaHArHAM/YiE=",
+      "_speculative_draft_primitives_cuda_6ee7cab.abi3.so": "a5VCjWMmiq92JDOwu/GDIZ+gVxgR7bjDHqSDIz9Bpgo=",
+      "_ops.py": "mXBQn2p+IHGD95tP5Z0LrYSEmTCuehbwhHIOcFElJg0=",
+      "speculative_draft_primitives/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  }
+}

build/torch211-cxx11-cu130-aarch64-linux/speculative_draft_primitives/__init__.py ADDED Viewed

	@@ -0,0 +1,14 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+def _import_from_path(file_path: Path):
+    path_hash = '{:x}'.format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    spec = importlib.util.spec_from_file_location(path_hash, file_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))