Promote latest kernel artifacts to main

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text
+*.pyd filter=lfs diff=lfs merge=lfs -text
+*.dylib filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,9 +0,0 @@
-# flashrt/speculative-draft-primitives
-
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
-
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/speculative-draft-primitives
-
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.

benchmarks/benchmark.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+"""Benchmark speculative-draft-primitives."""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+import sys
+
+import torch
+
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT / "speculative-draft-primitives" / "tests"))
+from test_speculative_draft_primitives import load_installed_ops, load_source_ops  # noqa: E402
+
+
+def time_us(fn, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    for _ in range(iters):
+        fn()
+    torch.cuda.synchronize()
+    return (time.perf_counter() - start) * 1e6 / iters
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["headline", "full"], default="headline")
+    parser.add_argument("--warmup", type=int, default=50)
+    parser.add_argument("--iters", type=int, default=200)
+    args = parser.parse_args()
+
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    shapes = [(16, 32000), (16, 248320)] if args.mode == "headline" else [
+        (1, 1024),
+        (4, 4096),
+        (16, 32000),
+        (16, 248320),
+    ]
+    print("rows,vocab,op,flashrt_us,torch_us,speedup,notes")
+    for rows, vocab in shapes:
+        logits = torch.randn((rows, vocab), device="cuda", dtype=torch.float32).to(torch.bfloat16)
+        drafts = torch.argmax(logits.float(), dim=1)[: min(rows, 15)].contiguous()
+        argmax_out = torch.empty((rows,), device=logits.device, dtype=torch.int64)
+        accept_n = torch.empty((1,), device=logits.device, dtype=torch.int32)
+        parts = 32 if vocab >= 131072 else (16 if vocab >= 65536 else (1 if vocab <= 4096 else 8))
+        partial_vals = torch.empty((rows, parts), device=logits.device, dtype=torch.float32)
+        partial_idx = torch.empty((rows, parts), device=logits.device, dtype=torch.int32)
+
+        if hasattr(ops, "ops"):
+            raw = ops.ops
+
+            def flash_argmax():
+                raw.argmax_bf16(logits, argmax_out)
+
+            def flash_accept_partitioned():
+                raw.accept_partitioned_bf16(
+                    logits, drafts, argmax_out, accept_n, partial_vals, partial_idx, min(rows, 15), parts
+                )
+        else:
+            def flash_argmax():
+                ops.argmax_bf16(logits, out=argmax_out)
+
+            def flash_accept_partitioned():
+                ops.accept_partitioned_bf16(
+                    logits,
+                    drafts,
+                    min(rows, 15),
+                    parts,
+                    argmax_out=argmax_out,
+                    accept_n=accept_n,
+                    partial_vals=partial_vals,
+                    partial_idx=partial_idx,
+                )
+
+        flash_us = time_us(flash_argmax, args.warmup, args.iters)
+        torch_us = time_us(lambda: torch.argmax(logits.float(), dim=1), args.warmup, args.iters)
+        print(f"{rows},{vocab},argmax_bf16,{flash_us:.3f},{torch_us:.3f},{torch_us / flash_us:.2f},static-output")
+        flash_us = time_us(flash_accept_partitioned, args.warmup, args.iters)
+        print(f"{rows},{vocab},accept_partitioned_bf16,{flash_us:.3f},n/a,n/a,static-workspace parts={parts}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

build/torch211-cxx11-cu128-x86_64-linux/__init__.py

@@ -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-cu128-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+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):
+    """
+    Prefix op by namespace.
+    """
+    return f"_speculative_draft_primitives_cuda_6ee7cab::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/_speculative_draft_primitives_cuda_6ee7cab.abi3.so

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

build/torch211-cxx11-cu128-x86_64-linux/metadata.json

@@ -0,0 +1,23 @@
+{
+  "name": "speculative-draft-primitives",
+  "id": "_speculative_draft_primitives_cuda_6ee7cab",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "10.1",
+      "12.0+PTX",
+      "7.0",
+      "7.2",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch211-cxx11-cu128-x86_64-linux/speculative_draft_primitives/__init__.py

@@ -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/torch211-cxx11-cu130-x86_64-linux/__init__.py

@@ -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-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+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):
+    """
+    Prefix op by namespace.
+    """
+    return f"_speculative_draft_primitives_cuda_6ee7cab::{op_name}"

build/torch211-cxx11-cu130-x86_64-linux/_speculative_draft_primitives_cuda_6ee7cab.abi3.so

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

build/torch211-cxx11-cu130-x86_64-linux/metadata.json

@@ -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": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch211-cxx11-cu130-x86_64-linux/speculative_draft_primitives/__init__.py

@@ -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-cu130-x86_64-linux/__init__.py

@@ -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/torch212-cxx11-cu130-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+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):
+    """
+    Prefix op by namespace.
+    """
+    return f"_speculative_draft_primitives_cuda_6ee7cab::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/_speculative_draft_primitives_cuda_6ee7cab.abi3.so

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

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

@@ -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": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch212-cxx11-cu130-x86_64-linux/speculative_draft_primitives/__init__.py

@@ -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-cu132-x86_64-linux/__init__.py

@@ -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/torch212-cxx11-cu132-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+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):
+    """
+    Prefix op by namespace.
+    """
+    return f"_speculative_draft_primitives_cuda_6ee7cab::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/_speculative_draft_primitives_cuda_6ee7cab.abi3.so

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

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

@@ -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": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch212-cxx11-cu132-x86_64-linux/speculative_draft_primitives/__init__.py

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