Promote latest kernel artifacts to main

README.md

@@ -1,9 +0,0 @@
-# flashrt/turboquant-kv
-
-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/turboquant-kv
-
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.

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

@@ -0,0 +1,169 @@
+"""FlashRT TurboQuant KV cache kernels."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def _check_packed(k_idx: torch.Tensor, k_qjl: torch.Tensor, v_idx: torch.Tensor) -> int:
+    if k_idx.dim() != 2 or k_idx.shape[1] != 128:
+        raise RuntimeError("k_idx_packed must have shape (M, 128)")
+    m = k_idx.shape[0]
+    if k_qjl.shape != (m, 32):
+        raise RuntimeError("k_qjl_packed must have shape (M, 32)")
+    if v_idx.shape != (m, 128):
+        raise RuntimeError("v_idx_packed must have shape (M, 128)")
+    return m
+
+
+@torch.library.register_fake(add_op_namespace_prefix("unpack_packed_bf16"))
+def _unpack_packed_bf16_fake(
+    k_idx_packed: torch.Tensor,
+    k_qjl_packed: torch.Tensor,
+    v_idx_packed: torch.Tensor,
+    cb_k_mse: torch.Tensor,
+    cb_v: torch.Tensor,
+    b_k_mse: int,
+    b_v: int,
+    y_k: torch.Tensor,
+    qjl_bf: torch.Tensor,
+    y_v: torch.Tensor,
+) -> None:
+    m = _check_packed(k_idx_packed, k_qjl_packed, v_idx_packed)
+    if y_k.shape != (m, 256) or qjl_bf.shape != (m, 256) or y_v.shape != (m, 256):
+        raise RuntimeError("outputs must have shape (M, 256)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("unpack_packed_mixed"))
+def _unpack_packed_mixed_fake(
+    k_idx_packed: torch.Tensor,
+    k_qjl_packed: torch.Tensor,
+    v_idx_packed: torch.Tensor,
+    cb_k_mse: torch.Tensor,
+    cb_v: torch.Tensor,
+    b_k_mse: int,
+    b_v: int,
+    y_k: torch.Tensor,
+    qjl_f: torch.Tensor,
+    y_v: torch.Tensor,
+) -> None:
+    m = _check_packed(k_idx_packed, k_qjl_packed, v_idx_packed)
+    if y_k.shape != (m, 256) or qjl_f.shape != (m, 256) or y_v.shape != (m, 256):
+        raise RuntimeError("outputs must have shape (M, 256)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("combine_kv_bf16"))
+def _combine_kv_bf16_fake(
+    k_mse: torch.Tensor,
+    k_qjl: torch.Tensor,
+    v_unit: torch.Tensor,
+    k_norm: torch.Tensor,
+    k_rnorm: torch.Tensor,
+    v_norm: torch.Tensor,
+    coef: float,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+) -> None:
+    if k_mse.dim() != 2 or k_mse.shape[1] != 256:
+        raise RuntimeError("k_mse must have shape (M, 256)")
+    m = k_mse.shape[0]
+    if k_qjl.shape != k_mse.shape or v_unit.shape != k_mse.shape:
+        raise RuntimeError("k_qjl and v_unit must have shape (M, 256)")
+    if k_norm.shape != (m,) or k_rnorm.shape != (m,) or v_norm.shape != (m,):
+        raise RuntimeError("norm tensors must have shape (M,)")
+    if k_out.shape != k_mse.shape or v_out.shape != k_mse.shape:
+        raise RuntimeError("outputs must have shape (M, 256)")
+    return None
+
+
+def unpack_packed_bf16(
+    k_idx_packed: torch.Tensor,
+    k_qjl_packed: torch.Tensor,
+    v_idx_packed: torch.Tensor,
+    cb_k_mse: torch.Tensor,
+    cb_v: torch.Tensor,
+    b_k_mse: int,
+    b_v: int,
+    *,
+    y_k: Optional[torch.Tensor] = None,
+    qjl_bf: Optional[torch.Tensor] = None,
+    y_v: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Unpack TurboQuant packed indices to BF16 `(M, 256)` tensors."""
+
+    m = k_idx_packed.shape[0]
+    if y_k is None:
+        y_k = torch.empty((m, 256), device=k_idx_packed.device, dtype=torch.bfloat16)
+    if qjl_bf is None:
+        qjl_bf = torch.empty_like(y_k)
+    if y_v is None:
+        y_v = torch.empty_like(y_k)
+    ops.unpack_packed_bf16(
+        k_idx_packed, k_qjl_packed, v_idx_packed, cb_k_mse, cb_v,
+        int(b_k_mse), int(b_v), y_k, qjl_bf, y_v,
+    )
+    return y_k, qjl_bf, y_v
+
+
+def unpack_packed_mixed(
+    k_idx_packed: torch.Tensor,
+    k_qjl_packed: torch.Tensor,
+    v_idx_packed: torch.Tensor,
+    cb_k_mse: torch.Tensor,
+    cb_v: torch.Tensor,
+    b_k_mse: int,
+    b_v: int,
+    *,
+    y_k: Optional[torch.Tensor] = None,
+    qjl_f: Optional[torch.Tensor] = None,
+    y_v: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Unpack packed K/V to BF16 K/V and FP32 QJL signs."""
+
+    m = k_idx_packed.shape[0]
+    if y_k is None:
+        y_k = torch.empty((m, 256), device=k_idx_packed.device, dtype=torch.bfloat16)
+    if qjl_f is None:
+        qjl_f = torch.empty((m, 256), device=k_idx_packed.device, dtype=torch.float32)
+    if y_v is None:
+        y_v = torch.empty_like(y_k)
+    ops.unpack_packed_mixed(
+        k_idx_packed, k_qjl_packed, v_idx_packed, cb_k_mse, cb_v,
+        int(b_k_mse), int(b_v), y_k, qjl_f, y_v,
+    )
+    return y_k, qjl_f, y_v
+
+
+def combine_kv_bf16(
+    k_mse: torch.Tensor,
+    k_qjl: torch.Tensor,
+    v_unit: torch.Tensor,
+    k_norm: torch.Tensor,
+    k_rnorm: torch.Tensor,
+    v_norm: torch.Tensor,
+    coef: float,
+    *,
+    k_out: Optional[torch.Tensor] = None,
+    v_out: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Combine TurboQuant K/V GEMM outputs into BF16 K/V cache tensors."""
+
+    if k_out is None:
+        k_out = torch.empty_like(k_mse)
+    if v_out is None:
+        v_out = torch.empty_like(v_unit)
+    ops.combine_kv_bf16(k_mse, k_qjl, v_unit, k_norm, k_rnorm, v_norm, float(coef), k_out, v_out)
+    return k_out, v_out
+
+
+__all__ = [
+    "combine_kv_bf16",
+    "unpack_packed_bf16",
+    "unpack_packed_mixed",
+]

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

@@ -0,0 +1,6 @@
+import torch
+from . import _turboquant_kv_cuda_5596053
+ops = torch.ops._turboquant_kv_cuda_5596053
+
+def add_op_namespace_prefix(op_name: str):
+    return f"_turboquant_kv_cuda_5596053::{op_name}"

build/torch211-cxx11-cu130-aarch64-linux/_turboquant_kv_cuda_5596053.abi3.so

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

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

@@ -0,0 +1,22 @@
+{
+  "name": "turboquant-kv",
+  "id": "_turboquant_kv_cuda_5596053",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "MRfXnYiu3ykN9lvSL/DKIK2TADkUKewmrj0CQynXZ0Y=",
+      "_turboquant_kv_cuda_5596053.abi3.so": "7RXZPOQ6wASAznmfVszo4xf2STx6jA1io4HJE1vk2E8=",
+      "_ops.py": "CB2L++MoBdyPQaPTodj5z5qaB08olDjuhe9ZAiApjG0=",
+      "turboquant_kv/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  }
+}

build/torch211-cxx11-cu130-aarch64-linux/turboquant_kv/__init__.py

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