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/fp8-kv-attention
-
-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/fp8-kv-attention
-
-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,97 @@
+#!/usr/bin/env python3
+"""Benchmark fp8-kv-attention against a PyTorch FP8-dequant reference."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import torch
+
+import sys
+
+TESTS = Path(__file__).resolve().parents[1] / "tests"
+sys.path.insert(0, str(TESTS))
+from test_fp8_kv_attention import SHAPES, SourceOps, load_installed_ops, load_source_ops, make_inputs, reference  # noqa: E402
+
+
+MODES = {
+    "smoke": ["decode_128"],
+    "headline": ["decode_1024", "verify4_1024", "verify8_4096"],
+    "full": ["decode_128", "decode_1024", "verify4_1024", "verify8_4096"],
+}
+
+
+def time_cuda(fn, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return float(start.elapsed_time(end) * 1000.0 / 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=sorted(MODES), default="smoke")
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iters", type=int, default=100)
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    rows = []
+    for name in MODES[args.mode]:
+        q_seq, kv_seq = SHAPES[name]
+        q, k, v = make_inputs(q_seq, kv_seq, seed=3000 + q_seq * 17 + kv_seq)
+        if isinstance(ops, SourceOps):
+            def kernel_call():
+                return ops.xqa_bf16_fp8kv(q, k, v, kv_seq)
+        else:
+            pages = k.shape[0]
+            page_table = ops.default_page_table(pages, device=q.device)
+            seq_lens = torch.tensor([[kv_seq]], device=q.device, dtype=torch.int32)
+            mask = ops.causal_spec_mask(q_seq, device=q.device)
+            sem, scratch = ops.allocate_workspace(q_seq=q_seq, device=q.device)
+            out = torch.empty_like(q)
+
+            def kernel_call():
+                return ops.xqa_bf16_fp8kv(
+                    q, k, v, page_table, seq_lens, mask,
+                    out=out, semaphores=sem, scratch=scratch,
+                    max_seq_len=pages * 128,
+                )
+
+        def ref_call():
+            return reference(q, k, v, kv_seq)
+
+        kernel_us = time_cuda(kernel_call, args.warmup, args.iters)
+        ref_us = time_cuda(ref_call, max(2, args.warmup // 5), max(5, args.iters // 10))
+        rows.append(
+            {
+                "shape": name,
+                "q_seq": q_seq,
+                "kv_seq": kv_seq,
+                "kernel_us": kernel_us,
+                "torch_reference_us": ref_us,
+                "speedup": ref_us / kernel_us,
+            }
+        )
+        print(f"{name}: kernel={kernel_us:.3f}us ref={ref_us:.3f}us speedup={ref_us / kernel_us:.2f}x")
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(json.dumps(rows, indent=2) + "\n")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

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

@@ -0,0 +1,168 @@
+"""FlashRT BF16-Q + FP8-KV XQA attention kernels."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+PAGE_SIZE = 128
+NUM_Q_HEADS = 24
+NUM_KV_HEADS = 4
+HEAD_DIM = 256
+
+
+@torch.library.register_fake(add_op_namespace_prefix("xqa_bf16_fp8kv"))
+def _xqa_bf16_fp8kv_fake(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    mask: torch.Tensor,
+    out: torch.Tensor,
+    semaphores: torch.Tensor,
+    scratch: torch.Tensor,
+    max_seq_len: int = 0,
+    q_scale: float = 1.0,
+    kv_scale: float = 1.0,
+    enable_pdl: bool = True,
+    sm_count: int = 0,
+    k_stride_page: int = 0,
+    k_stride_token: int = 0,
+    k_stride_head: int = 0,
+) -> None:
+    if q.dim() == 3:
+        q_seq = q.shape[0]
+        ok = q.shape[1:] == (NUM_Q_HEADS, HEAD_DIM)
+    elif q.dim() == 5:
+        q_seq = q.shape[2]
+        ok = q.shape[:2] == (1, 1) and q.shape[3:] == (NUM_Q_HEADS, HEAD_DIM)
+    else:
+        raise RuntimeError("q must have rank 3 or 5")
+    if not ok or out.shape != q.shape:
+        raise RuntimeError("q/out shape mismatch for v1 XQA contract")
+    if k_cache.dim() != 4 or k_cache.shape[1:] != (PAGE_SIZE, NUM_KV_HEADS, HEAD_DIM):
+        raise RuntimeError("k_cache must have shape (pages,128,4,256)")
+    if v_cache.shape != k_cache.shape:
+        raise RuntimeError("v_cache shape mismatch")
+    if mask.numel() < q_seq * ((q_seq + 31) // 32):
+        raise RuntimeError("mask is too small")
+    return None
+
+
+def causal_spec_mask(q_seq: int, *, device: torch.device | str = "cuda", dtype: torch.dtype = torch.int32) -> torch.Tensor:
+    """Return the packed lower-triangular mask expected by the v1 XQA kernel."""
+
+    q_seq = int(q_seq)
+    words = (q_seq + 31) // 32
+    rows = torch.zeros((q_seq, words), dtype=torch.int32)
+    for i in range(q_seq):
+        upto = i + 1
+        full = upto // 32
+        rem = upto % 32
+        if full:
+            rows[i, :full] = -1
+        if rem:
+            rows[i, full] = (1 << rem) - 1
+    return rows.to(device=device, dtype=dtype)
+
+
+def default_page_table(num_pages: int, *, device: torch.device | str = "cuda") -> torch.Tensor:
+    """Contiguous one-batch page table for `(pages,128,4,256)` K/V caches."""
+
+    return torch.arange(int(num_pages), device=device, dtype=torch.int32).view(1, int(num_pages))
+
+
+def allocate_workspace(
+    *,
+    q_seq: int,
+    device: torch.device | str = "cuda",
+    scratch_mb: int = 256,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Allocate semaphores and scratch tensors for static-buffer runtimes."""
+
+    sem_count = NUM_KV_HEADS * (((int(q_seq) * (NUM_Q_HEADS // NUM_KV_HEADS)) + 31) // 32)
+    semaphores = torch.zeros(max(256, sem_count), device=device, dtype=torch.int32)
+    scratch = torch.empty(max(1, int(scratch_mb)) << 20, device=device, dtype=torch.uint8)
+    return semaphores, scratch
+
+
+def xqa_bf16_fp8kv(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_table: Optional[torch.Tensor] = None,
+    seq_lens: Optional[torch.Tensor] = None,
+    mask: Optional[torch.Tensor] = None,
+    *,
+    out: Optional[torch.Tensor] = None,
+    semaphores: Optional[torch.Tensor] = None,
+    scratch: Optional[torch.Tensor] = None,
+    max_seq_len: int = 0,
+    q_scale: float = 1.0,
+    kv_scale: float = 1.0,
+    enable_pdl: bool = True,
+    sm_count: int = 0,
+    k_stride_page: int = 0,
+    k_stride_token: int = 0,
+    k_stride_head: int = 0,
+) -> torch.Tensor:
+    """Run BF16-query / FP8-KV XQA attention for the v1 fixed public shape.
+
+    v1 shape contract:
+    `q`: `(q_seq, 24, 256)` or `(1, 1, q_seq, 24, 256)` BF16.
+    `k_cache`, `v_cache`: `(pages, 128, 4, 256)` FP8 E4M3.
+    """
+
+    if out is None:
+        out = torch.empty_like(q)
+    if page_table is None:
+        page_table = default_page_table(k_cache.shape[0], device=q.device)
+    if seq_lens is None:
+        seq_lens = torch.tensor([[k_cache.shape[0] * PAGE_SIZE]], device=q.device, dtype=torch.int32)
+    if mask is None:
+        q_seq = q.shape[0] if q.dim() == 3 else q.shape[2]
+        mask = causal_spec_mask(int(q_seq), device=q.device, dtype=torch.int32)
+    if semaphores is None or scratch is None:
+        q_seq = q.shape[0] if q.dim() == 3 else q.shape[2]
+        semaphores_new, scratch_new = allocate_workspace(q_seq=int(q_seq), device=q.device)
+        if semaphores is None:
+            semaphores = semaphores_new
+        if scratch is None:
+            scratch = scratch_new
+    ops.xqa_bf16_fp8kv(
+        q,
+        k_cache,
+        v_cache,
+        page_table,
+        seq_lens,
+        mask,
+        out,
+        semaphores,
+        scratch,
+        int(max_seq_len),
+        float(q_scale),
+        float(kv_scale),
+        bool(enable_pdl),
+        int(sm_count),
+        int(k_stride_page),
+        int(k_stride_token),
+        int(k_stride_head),
+    )
+    return out
+
+
+__all__ = [
+    "HEAD_DIM",
+    "NUM_KV_HEADS",
+    "NUM_Q_HEADS",
+    "PAGE_SIZE",
+    "allocate_workspace",
+    "causal_spec_mask",
+    "default_page_table",
+    "xqa_bf16_fp8kv",
+]

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

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

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

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

build/torch211-cxx11-cu128-x86_64-linux/fp8_kv_attention/__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-cu128-x86_64-linux/metadata.json

@@ -0,0 +1,22 @@
+{
+  "name": "fp8-kv-attention",
+  "id": "_fp8_kv_attention_cuda_1798e7f",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "AipNkYPvsd/zyT+YU7bggBObENJYocXbAV7/Yr5A+6c=",
+      "_fp8_kv_attention_cuda_1798e7f.abi3.so": "z0AaImjYf/6DbYJMOfsyhvoSgo66ZKXjwn5hhFVW/NM=",
+      "_ops.py": "u1tuL5lFum0BIgl8+j1z86OkjJn3e4Mnmc/29zQSfHA=",
+      "fp8_kv_attention/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

@@ -0,0 +1,168 @@
+"""FlashRT BF16-Q + FP8-KV XQA attention kernels."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+PAGE_SIZE = 128
+NUM_Q_HEADS = 24
+NUM_KV_HEADS = 4
+HEAD_DIM = 256
+
+
+@torch.library.register_fake(add_op_namespace_prefix("xqa_bf16_fp8kv"))
+def _xqa_bf16_fp8kv_fake(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    mask: torch.Tensor,
+    out: torch.Tensor,
+    semaphores: torch.Tensor,
+    scratch: torch.Tensor,
+    max_seq_len: int = 0,
+    q_scale: float = 1.0,
+    kv_scale: float = 1.0,
+    enable_pdl: bool = True,
+    sm_count: int = 0,
+    k_stride_page: int = 0,
+    k_stride_token: int = 0,
+    k_stride_head: int = 0,
+) -> None:
+    if q.dim() == 3:
+        q_seq = q.shape[0]
+        ok = q.shape[1:] == (NUM_Q_HEADS, HEAD_DIM)
+    elif q.dim() == 5:
+        q_seq = q.shape[2]
+        ok = q.shape[:2] == (1, 1) and q.shape[3:] == (NUM_Q_HEADS, HEAD_DIM)
+    else:
+        raise RuntimeError("q must have rank 3 or 5")
+    if not ok or out.shape != q.shape:
+        raise RuntimeError("q/out shape mismatch for v1 XQA contract")
+    if k_cache.dim() != 4 or k_cache.shape[1:] != (PAGE_SIZE, NUM_KV_HEADS, HEAD_DIM):
+        raise RuntimeError("k_cache must have shape (pages,128,4,256)")
+    if v_cache.shape != k_cache.shape:
+        raise RuntimeError("v_cache shape mismatch")
+    if mask.numel() < q_seq * ((q_seq + 31) // 32):
+        raise RuntimeError("mask is too small")
+    return None
+
+
+def causal_spec_mask(q_seq: int, *, device: torch.device | str = "cuda", dtype: torch.dtype = torch.int32) -> torch.Tensor:
+    """Return the packed lower-triangular mask expected by the v1 XQA kernel."""
+
+    q_seq = int(q_seq)
+    words = (q_seq + 31) // 32
+    rows = torch.zeros((q_seq, words), dtype=torch.int32)
+    for i in range(q_seq):
+        upto = i + 1
+        full = upto // 32
+        rem = upto % 32
+        if full:
+            rows[i, :full] = -1
+        if rem:
+            rows[i, full] = (1 << rem) - 1
+    return rows.to(device=device, dtype=dtype)
+
+
+def default_page_table(num_pages: int, *, device: torch.device | str = "cuda") -> torch.Tensor:
+    """Contiguous one-batch page table for `(pages,128,4,256)` K/V caches."""
+
+    return torch.arange(int(num_pages), device=device, dtype=torch.int32).view(1, int(num_pages))
+
+
+def allocate_workspace(
+    *,
+    q_seq: int,
+    device: torch.device | str = "cuda",
+    scratch_mb: int = 256,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Allocate semaphores and scratch tensors for static-buffer runtimes."""
+
+    sem_count = NUM_KV_HEADS * (((int(q_seq) * (NUM_Q_HEADS // NUM_KV_HEADS)) + 31) // 32)
+    semaphores = torch.zeros(max(256, sem_count), device=device, dtype=torch.int32)
+    scratch = torch.empty(max(1, int(scratch_mb)) << 20, device=device, dtype=torch.uint8)
+    return semaphores, scratch
+
+
+def xqa_bf16_fp8kv(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_table: Optional[torch.Tensor] = None,
+    seq_lens: Optional[torch.Tensor] = None,
+    mask: Optional[torch.Tensor] = None,
+    *,
+    out: Optional[torch.Tensor] = None,
+    semaphores: Optional[torch.Tensor] = None,
+    scratch: Optional[torch.Tensor] = None,
+    max_seq_len: int = 0,
+    q_scale: float = 1.0,
+    kv_scale: float = 1.0,
+    enable_pdl: bool = True,
+    sm_count: int = 0,
+    k_stride_page: int = 0,
+    k_stride_token: int = 0,
+    k_stride_head: int = 0,
+) -> torch.Tensor:
+    """Run BF16-query / FP8-KV XQA attention for the v1 fixed public shape.
+
+    v1 shape contract:
+    `q`: `(q_seq, 24, 256)` or `(1, 1, q_seq, 24, 256)` BF16.
+    `k_cache`, `v_cache`: `(pages, 128, 4, 256)` FP8 E4M3.
+    """
+
+    if out is None:
+        out = torch.empty_like(q)
+    if page_table is None:
+        page_table = default_page_table(k_cache.shape[0], device=q.device)
+    if seq_lens is None:
+        seq_lens = torch.tensor([[k_cache.shape[0] * PAGE_SIZE]], device=q.device, dtype=torch.int32)
+    if mask is None:
+        q_seq = q.shape[0] if q.dim() == 3 else q.shape[2]
+        mask = causal_spec_mask(int(q_seq), device=q.device, dtype=torch.int32)
+    if semaphores is None or scratch is None:
+        q_seq = q.shape[0] if q.dim() == 3 else q.shape[2]
+        semaphores_new, scratch_new = allocate_workspace(q_seq=int(q_seq), device=q.device)
+        if semaphores is None:
+            semaphores = semaphores_new
+        if scratch is None:
+            scratch = scratch_new
+    ops.xqa_bf16_fp8kv(
+        q,
+        k_cache,
+        v_cache,
+        page_table,
+        seq_lens,
+        mask,
+        out,
+        semaphores,
+        scratch,
+        int(max_seq_len),
+        float(q_scale),
+        float(kv_scale),
+        bool(enable_pdl),
+        int(sm_count),
+        int(k_stride_page),
+        int(k_stride_token),
+        int(k_stride_head),
+    )
+    return out
+
+
+__all__ = [
+    "HEAD_DIM",
+    "NUM_KV_HEADS",
+    "NUM_Q_HEADS",
+    "PAGE_SIZE",
+    "allocate_workspace",
+    "causal_spec_mask",
+    "default_page_table",
+    "xqa_bf16_fp8kv",
+]

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

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

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

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

build/torch211-cxx11-cu130-x86_64-linux/fp8_kv_attention/__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/metadata.json

@@ -0,0 +1,22 @@
+{
+  "name": "fp8-kv-attention",
+  "id": "_fp8_kv_attention_cuda_1798e7f",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "AipNkYPvsd/zyT+YU7bggBObENJYocXbAV7/Yr5A+6c=",
+      "_fp8_kv_attention_cuda_1798e7f.abi3.so": "7lNvL0b1EY0JXBSiLrvydYEaOcOrpWt3Shxnl8I0BEo=",
+      "_ops.py": "u1tuL5lFum0BIgl8+j1z86OkjJn3e4Mnmc/29zQSfHA=",
+      "fp8_kv_attention/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

@@ -0,0 +1,168 @@
+"""FlashRT BF16-Q + FP8-KV XQA attention kernels."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+PAGE_SIZE = 128
+NUM_Q_HEADS = 24
+NUM_KV_HEADS = 4
+HEAD_DIM = 256
+
+
+@torch.library.register_fake(add_op_namespace_prefix("xqa_bf16_fp8kv"))
+def _xqa_bf16_fp8kv_fake(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    mask: torch.Tensor,
+    out: torch.Tensor,
+    semaphores: torch.Tensor,
+    scratch: torch.Tensor,
+    max_seq_len: int = 0,
+    q_scale: float = 1.0,
+    kv_scale: float = 1.0,
+    enable_pdl: bool = True,
+    sm_count: int = 0,
+    k_stride_page: int = 0,
+    k_stride_token: int = 0,
+    k_stride_head: int = 0,
+) -> None:
+    if q.dim() == 3:
+        q_seq = q.shape[0]
+        ok = q.shape[1:] == (NUM_Q_HEADS, HEAD_DIM)
+    elif q.dim() == 5:
+        q_seq = q.shape[2]
+        ok = q.shape[:2] == (1, 1) and q.shape[3:] == (NUM_Q_HEADS, HEAD_DIM)
+    else:
+        raise RuntimeError("q must have rank 3 or 5")
+    if not ok or out.shape != q.shape:
+        raise RuntimeError("q/out shape mismatch for v1 XQA contract")
+    if k_cache.dim() != 4 or k_cache.shape[1:] != (PAGE_SIZE, NUM_KV_HEADS, HEAD_DIM):
+        raise RuntimeError("k_cache must have shape (pages,128,4,256)")
+    if v_cache.shape != k_cache.shape:
+        raise RuntimeError("v_cache shape mismatch")
+    if mask.numel() < q_seq * ((q_seq + 31) // 32):
+        raise RuntimeError("mask is too small")
+    return None
+
+
+def causal_spec_mask(q_seq: int, *, device: torch.device | str = "cuda", dtype: torch.dtype = torch.int32) -> torch.Tensor:
+    """Return the packed lower-triangular mask expected by the v1 XQA kernel."""
+
+    q_seq = int(q_seq)
+    words = (q_seq + 31) // 32
+    rows = torch.zeros((q_seq, words), dtype=torch.int32)
+    for i in range(q_seq):
+        upto = i + 1
+        full = upto // 32
+        rem = upto % 32
+        if full:
+            rows[i, :full] = -1
+        if rem:
+            rows[i, full] = (1 << rem) - 1
+    return rows.to(device=device, dtype=dtype)
+
+
+def default_page_table(num_pages: int, *, device: torch.device | str = "cuda") -> torch.Tensor:
+    """Contiguous one-batch page table for `(pages,128,4,256)` K/V caches."""
+
+    return torch.arange(int(num_pages), device=device, dtype=torch.int32).view(1, int(num_pages))
+
+
+def allocate_workspace(
+    *,
+    q_seq: int,
+    device: torch.device | str = "cuda",
+    scratch_mb: int = 256,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Allocate semaphores and scratch tensors for static-buffer runtimes."""
+
+    sem_count = NUM_KV_HEADS * (((int(q_seq) * (NUM_Q_HEADS // NUM_KV_HEADS)) + 31) // 32)
+    semaphores = torch.zeros(max(256, sem_count), device=device, dtype=torch.int32)
+    scratch = torch.empty(max(1, int(scratch_mb)) << 20, device=device, dtype=torch.uint8)
+    return semaphores, scratch
+
+
+def xqa_bf16_fp8kv(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_table: Optional[torch.Tensor] = None,
+    seq_lens: Optional[torch.Tensor] = None,
+    mask: Optional[torch.Tensor] = None,
+    *,
+    out: Optional[torch.Tensor] = None,
+    semaphores: Optional[torch.Tensor] = None,
+    scratch: Optional[torch.Tensor] = None,
+    max_seq_len: int = 0,
+    q_scale: float = 1.0,
+    kv_scale: float = 1.0,
+    enable_pdl: bool = True,
+    sm_count: int = 0,
+    k_stride_page: int = 0,
+    k_stride_token: int = 0,
+    k_stride_head: int = 0,
+) -> torch.Tensor:
+    """Run BF16-query / FP8-KV XQA attention for the v1 fixed public shape.
+
+    v1 shape contract:
+    `q`: `(q_seq, 24, 256)` or `(1, 1, q_seq, 24, 256)` BF16.
+    `k_cache`, `v_cache`: `(pages, 128, 4, 256)` FP8 E4M3.
+    """
+
+    if out is None:
+        out = torch.empty_like(q)
+    if page_table is None:
+        page_table = default_page_table(k_cache.shape[0], device=q.device)
+    if seq_lens is None:
+        seq_lens = torch.tensor([[k_cache.shape[0] * PAGE_SIZE]], device=q.device, dtype=torch.int32)
+    if mask is None:
+        q_seq = q.shape[0] if q.dim() == 3 else q.shape[2]
+        mask = causal_spec_mask(int(q_seq), device=q.device, dtype=torch.int32)
+    if semaphores is None or scratch is None:
+        q_seq = q.shape[0] if q.dim() == 3 else q.shape[2]
+        semaphores_new, scratch_new = allocate_workspace(q_seq=int(q_seq), device=q.device)
+        if semaphores is None:
+            semaphores = semaphores_new
+        if scratch is None:
+            scratch = scratch_new
+    ops.xqa_bf16_fp8kv(
+        q,
+        k_cache,
+        v_cache,
+        page_table,
+        seq_lens,
+        mask,
+        out,
+        semaphores,
+        scratch,
+        int(max_seq_len),
+        float(q_scale),
+        float(kv_scale),
+        bool(enable_pdl),
+        int(sm_count),
+        int(k_stride_page),
+        int(k_stride_token),
+        int(k_stride_head),
+    )
+    return out
+
+
+__all__ = [
+    "HEAD_DIM",
+    "NUM_KV_HEADS",
+    "NUM_Q_HEADS",
+    "PAGE_SIZE",
+    "allocate_workspace",
+    "causal_spec_mask",
+    "default_page_table",
+    "xqa_bf16_fp8kv",
+]

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

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

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

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

build/torch212-cxx11-cu130-x86_64-linux/fp8_kv_attention/__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/metadata.json

@@ -0,0 +1,22 @@
+{
+  "name": "fp8-kv-attention",
+  "id": "_fp8_kv_attention_cuda_1798e7f",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "AipNkYPvsd/zyT+YU7bggBObENJYocXbAV7/Yr5A+6c=",
+      "_fp8_kv_attention_cuda_1798e7f.abi3.so": "fzVV8jnbvBRLRjyakPmR+Cq9sxg442AjZpfIWbReHMA=",
+      "_ops.py": "u1tuL5lFum0BIgl8+j1z86OkjJn3e4Mnmc/29zQSfHA=",
+      "fp8_kv_attention/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

@@ -0,0 +1,168 @@
+"""FlashRT BF16-Q + FP8-KV XQA attention kernels."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+PAGE_SIZE = 128
+NUM_Q_HEADS = 24
+NUM_KV_HEADS = 4
+HEAD_DIM = 256
+
+
+@torch.library.register_fake(add_op_namespace_prefix("xqa_bf16_fp8kv"))
+def _xqa_bf16_fp8kv_fake(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_table: torch.Tensor,
+    seq_lens: torch.Tensor,
+    mask: torch.Tensor,
+    out: torch.Tensor,
+    semaphores: torch.Tensor,
+    scratch: torch.Tensor,
+    max_seq_len: int = 0,
+    q_scale: float = 1.0,
+    kv_scale: float = 1.0,
+    enable_pdl: bool = True,
+    sm_count: int = 0,
+    k_stride_page: int = 0,
+    k_stride_token: int = 0,
+    k_stride_head: int = 0,
+) -> None:
+    if q.dim() == 3:
+        q_seq = q.shape[0]
+        ok = q.shape[1:] == (NUM_Q_HEADS, HEAD_DIM)
+    elif q.dim() == 5:
+        q_seq = q.shape[2]
+        ok = q.shape[:2] == (1, 1) and q.shape[3:] == (NUM_Q_HEADS, HEAD_DIM)
+    else:
+        raise RuntimeError("q must have rank 3 or 5")
+    if not ok or out.shape != q.shape:
+        raise RuntimeError("q/out shape mismatch for v1 XQA contract")
+    if k_cache.dim() != 4 or k_cache.shape[1:] != (PAGE_SIZE, NUM_KV_HEADS, HEAD_DIM):
+        raise RuntimeError("k_cache must have shape (pages,128,4,256)")
+    if v_cache.shape != k_cache.shape:
+        raise RuntimeError("v_cache shape mismatch")
+    if mask.numel() < q_seq * ((q_seq + 31) // 32):
+        raise RuntimeError("mask is too small")
+    return None
+
+
+def causal_spec_mask(q_seq: int, *, device: torch.device | str = "cuda", dtype: torch.dtype = torch.int32) -> torch.Tensor:
+    """Return the packed lower-triangular mask expected by the v1 XQA kernel."""
+
+    q_seq = int(q_seq)
+    words = (q_seq + 31) // 32
+    rows = torch.zeros((q_seq, words), dtype=torch.int32)
+    for i in range(q_seq):
+        upto = i + 1
+        full = upto // 32
+        rem = upto % 32
+        if full:
+            rows[i, :full] = -1
+        if rem:
+            rows[i, full] = (1 << rem) - 1
+    return rows.to(device=device, dtype=dtype)
+
+
+def default_page_table(num_pages: int, *, device: torch.device | str = "cuda") -> torch.Tensor:
+    """Contiguous one-batch page table for `(pages,128,4,256)` K/V caches."""
+
+    return torch.arange(int(num_pages), device=device, dtype=torch.int32).view(1, int(num_pages))
+
+
+def allocate_workspace(
+    *,
+    q_seq: int,
+    device: torch.device | str = "cuda",
+    scratch_mb: int = 256,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Allocate semaphores and scratch tensors for static-buffer runtimes."""
+
+    sem_count = NUM_KV_HEADS * (((int(q_seq) * (NUM_Q_HEADS // NUM_KV_HEADS)) + 31) // 32)
+    semaphores = torch.zeros(max(256, sem_count), device=device, dtype=torch.int32)
+    scratch = torch.empty(max(1, int(scratch_mb)) << 20, device=device, dtype=torch.uint8)
+    return semaphores, scratch
+
+
+def xqa_bf16_fp8kv(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_table: Optional[torch.Tensor] = None,
+    seq_lens: Optional[torch.Tensor] = None,
+    mask: Optional[torch.Tensor] = None,
+    *,
+    out: Optional[torch.Tensor] = None,
+    semaphores: Optional[torch.Tensor] = None,
+    scratch: Optional[torch.Tensor] = None,
+    max_seq_len: int = 0,
+    q_scale: float = 1.0,
+    kv_scale: float = 1.0,
+    enable_pdl: bool = True,
+    sm_count: int = 0,
+    k_stride_page: int = 0,
+    k_stride_token: int = 0,
+    k_stride_head: int = 0,
+) -> torch.Tensor:
+    """Run BF16-query / FP8-KV XQA attention for the v1 fixed public shape.
+
+    v1 shape contract:
+    `q`: `(q_seq, 24, 256)` or `(1, 1, q_seq, 24, 256)` BF16.
+    `k_cache`, `v_cache`: `(pages, 128, 4, 256)` FP8 E4M3.
+    """
+
+    if out is None:
+        out = torch.empty_like(q)
+    if page_table is None:
+        page_table = default_page_table(k_cache.shape[0], device=q.device)
+    if seq_lens is None:
+        seq_lens = torch.tensor([[k_cache.shape[0] * PAGE_SIZE]], device=q.device, dtype=torch.int32)
+    if mask is None:
+        q_seq = q.shape[0] if q.dim() == 3 else q.shape[2]
+        mask = causal_spec_mask(int(q_seq), device=q.device, dtype=torch.int32)
+    if semaphores is None or scratch is None:
+        q_seq = q.shape[0] if q.dim() == 3 else q.shape[2]
+        semaphores_new, scratch_new = allocate_workspace(q_seq=int(q_seq), device=q.device)
+        if semaphores is None:
+            semaphores = semaphores_new
+        if scratch is None:
+            scratch = scratch_new
+    ops.xqa_bf16_fp8kv(
+        q,
+        k_cache,
+        v_cache,
+        page_table,
+        seq_lens,
+        mask,
+        out,
+        semaphores,
+        scratch,
+        int(max_seq_len),
+        float(q_scale),
+        float(kv_scale),
+        bool(enable_pdl),
+        int(sm_count),
+        int(k_stride_page),
+        int(k_stride_token),
+        int(k_stride_head),
+    )
+    return out
+
+
+__all__ = [
+    "HEAD_DIM",
+    "NUM_KV_HEADS",
+    "NUM_Q_HEADS",
+    "PAGE_SIZE",
+    "allocate_workspace",
+    "causal_spec_mask",
+    "default_page_table",
+    "xqa_bf16_fp8kv",
+]

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

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

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

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

build/torch212-cxx11-cu132-x86_64-linux/fp8_kv_attention/__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/metadata.json

@@ -0,0 +1,22 @@
+{
+  "name": "fp8-kv-attention",
+  "id": "_fp8_kv_attention_cuda_1798e7f",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "AipNkYPvsd/zyT+YU7bggBObENJYocXbAV7/Yr5A+6c=",
+      "_fp8_kv_attention_cuda_1798e7f.abi3.so": "x+RxX1RWNYThnvPFB/fd19DtMpfii326L9H5vKy2gFE=",
+      "_ops.py": "u1tuL5lFum0BIgl8+j1z86OkjJn3e4Mnmc/29zQSfHA=",
+      "fp8_kv_attention/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}