Promote latest kernel artifacts to main

README.md

@@ -1,9 +1,47 @@
-# flashrt/fp8-kv-attention
+---
+license: apache-2.0
+tags:
+- cuda
+- flashrt
+- attention
+- fp8
+- kv-cache
+- transformers
+---
 
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
+# FlashRT FP8 KV Attention
 
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/fp8-kv-attention
+Native CUDA XQA attention for BF16 queries over FP8 E4M3 paged K/V cache.
 
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.
+## Available Functions
+
+- `xqa_bf16_fp8kv`
+- `causal_spec_mask`
+- `default_page_table`
+- `allocate_workspace`
+
+## Scope
+
+v2 supports validated production decode/verify profiles:
+
+- BF16 Q/O
+- FP8 E4M3 K/V cache
+- Q/KV/head-dim `24/4/256`, `32/8/128`, `32/16/128`, and `16/8/128`
+- page size `128`
+- speculative/decode `q_seq <= 32`
+
+This is not a generic FlashAttention replacement. It is the direct FP8-KV XQA
+decode/verify kernel used to keep long-context transformer runtimes off BF16
+KV cache bandwidth.
+
+## Minimal Usage
+
+```python
+from kernels import get_kernel
+
+attn = get_kernel("flashrt/fp8-kv-attention", version=2, trust_remote_code=True)
+out = attn.xqa_bf16_fp8kv(q_bf16, k_cache_fp8, v_cache_fp8)
+```
+
+Pass explicit `page_table`, `seq_lens`, `mask`, `out`, `semaphores`, and
+`scratch` tensors for CUDA Graph/static-buffer runtimes.

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

@@ -0,0 +1,182 @@
+"""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
+SUPPORTED_CONFIGS = (
+    (24, 4, 256),
+    (16, 2, 256),
+    (32, 8, 128),
+    (32, 16, 128),
+    (16, 8, 128),
+)
+
+
+@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]
+        q_heads, head_dim = q.shape[1:]
+    elif q.dim() == 5:
+        q_seq = q.shape[2]
+        if q.shape[:2] != (1, 1):
+            raise RuntimeError("rank-5 q must begin with (1,1)")
+        q_heads, head_dim = q.shape[3:]
+    else:
+        raise RuntimeError("q must have rank 3 or 5")
+    if out.shape != q.shape:
+        raise RuntimeError("q/out shape mismatch")
+    if k_cache.dim() != 4 or k_cache.shape[1] != PAGE_SIZE or k_cache.shape[3] != head_dim:
+        raise RuntimeError("k_cache shape does not match q")
+    if v_cache.shape != k_cache.shape:
+        raise RuntimeError("v_cache shape mismatch")
+    config = (q_heads, k_cache.shape[2], head_dim)
+    if config not in SUPPORTED_CONFIGS:
+        raise RuntimeError(f"unsupported XQA configuration: {config}")
+    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 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:
+    """Return a contiguous one-batch page table."""
+
+    return torch.arange(int(num_pages), device=device, dtype=torch.int32).view(1, int(num_pages))
+
+
+def allocate_workspace(
+    *,
+    q_seq: int,
+    num_q_heads: int = 24,
+    num_kv_heads: int = 4,
+    device: torch.device | str = "cuda",
+    scratch_mb: int = 256,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Allocate semaphores and scratch tensors for static-buffer runtimes."""
+
+    if num_q_heads % num_kv_heads:
+        raise ValueError("num_q_heads must be divisible by num_kv_heads")
+    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 a validated BF16-query / FP8-KV paged XQA configuration."""
+
+    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]
+        q_heads = q.shape[1] if q.dim() == 3 else q.shape[3]
+        semaphores_new, scratch_new = allocate_workspace(
+            q_seq=int(q_seq),
+            num_q_heads=int(q_heads),
+            num_kv_heads=int(k_cache.shape[2]),
+            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__ = [
+    "PAGE_SIZE",
+    "SUPPORTED_CONFIGS",
+    "allocate_workspace",
+    "causal_spec_mask",
+    "default_page_table",
+    "xqa_bf16_fp8kv",
+]

build/torch213-cxx11-cu130-x86_64-linux/_fp8_kv_attention_cuda_e18d34b.abi3.so

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

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

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

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

@@ -0,0 +1,34 @@
+{
+  "name": "fp8-kv-attention",
+  "id": "_fp8_kv_attention_cuda_e18d34b",
+  "version": 3,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0",
+      "12.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "+u58XepDIAhYF7g3KZSoRhkN/PInROHfZilcN7Th+nc=",
+      "_fp8_kv_attention_cuda_e18d34b.abi3.so": "quhdZ4OU1f3d/3CAZQ/RWVQN9y59fCOev3wzrbuQQyo=",
+      "_ops.py": "RDYB4/Nm3ZFmR0Z0Hooh/I/f8g7lJ02RDZxkt+SUp7E=",
+      "fp8_kv_attention/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "b39ca23f1b36383df00b27b3ffe1276cd5dbea85",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "e18d34b116c127d0c976303a3e05c78ec47760d1",
+      "dirty": false
+    }
+  }
+}

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

@@ -0,0 +1,182 @@
+"""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
+SUPPORTED_CONFIGS = (
+    (24, 4, 256),
+    (16, 2, 256),
+    (32, 8, 128),
+    (32, 16, 128),
+    (16, 8, 128),
+)
+
+
+@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]
+        q_heads, head_dim = q.shape[1:]
+    elif q.dim() == 5:
+        q_seq = q.shape[2]
+        if q.shape[:2] != (1, 1):
+            raise RuntimeError("rank-5 q must begin with (1,1)")
+        q_heads, head_dim = q.shape[3:]
+    else:
+        raise RuntimeError("q must have rank 3 or 5")
+    if out.shape != q.shape:
+        raise RuntimeError("q/out shape mismatch")
+    if k_cache.dim() != 4 or k_cache.shape[1] != PAGE_SIZE or k_cache.shape[3] != head_dim:
+        raise RuntimeError("k_cache shape does not match q")
+    if v_cache.shape != k_cache.shape:
+        raise RuntimeError("v_cache shape mismatch")
+    config = (q_heads, k_cache.shape[2], head_dim)
+    if config not in SUPPORTED_CONFIGS:
+        raise RuntimeError(f"unsupported XQA configuration: {config}")
+    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 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:
+    """Return a contiguous one-batch page table."""
+
+    return torch.arange(int(num_pages), device=device, dtype=torch.int32).view(1, int(num_pages))
+
+
+def allocate_workspace(
+    *,
+    q_seq: int,
+    num_q_heads: int = 24,
+    num_kv_heads: int = 4,
+    device: torch.device | str = "cuda",
+    scratch_mb: int = 256,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Allocate semaphores and scratch tensors for static-buffer runtimes."""
+
+    if num_q_heads % num_kv_heads:
+        raise ValueError("num_q_heads must be divisible by num_kv_heads")
+    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 a validated BF16-query / FP8-KV paged XQA configuration."""
+
+    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]
+        q_heads = q.shape[1] if q.dim() == 3 else q.shape[3]
+        semaphores_new, scratch_new = allocate_workspace(
+            q_seq=int(q_seq),
+            num_q_heads=int(q_heads),
+            num_kv_heads=int(k_cache.shape[2]),
+            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__ = [
+    "PAGE_SIZE",
+    "SUPPORTED_CONFIGS",
+    "allocate_workspace",
+    "causal_spec_mask",
+    "default_page_table",
+    "xqa_bf16_fp8kv",
+]

build/torch213-cxx11-cu132-x86_64-linux/_fp8_kv_attention_cuda_e18d34b.abi3.so

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

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

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

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

@@ -0,0 +1,34 @@
+{
+  "name": "fp8-kv-attention",
+  "id": "_fp8_kv_attention_cuda_e18d34b",
+  "version": 3,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0",
+      "12.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "+u58XepDIAhYF7g3KZSoRhkN/PInROHfZilcN7Th+nc=",
+      "_fp8_kv_attention_cuda_e18d34b.abi3.so": "jESvmZOW5jp3JSOLacn/Wzv5ypgYyfRtTnxKrEsWwns=",
+      "_ops.py": "RDYB4/Nm3ZFmR0Z0Hooh/I/f8g7lJ02RDZxkt+SUp7E=",
+      "fp8_kv_attention/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "b39ca23f1b36383df00b27b3ffe1276cd5dbea85",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "e18d34b116c127d0c976303a3e05c78ec47760d1",
+      "dirty": false
+    }
+  }
+}