Uploaded using `kernel-builder`.

build/torch210-cxx11-cu126-x86_64-linux/_ops.py

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

build/torch210-cxx11-cu126-x86_64-linux/{_paged_attention_cuda_86f75d9.abi3.so → _paged_attention_cuda_2567cd6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f3b914c6cb87a56bf56630bb51a53ec05762bbbcca7b61bd8cdb1b07e1b4978c
+oid sha256:78a5a91ffa5a61aa74fb403262b0558f223898b63777968d8efbf6f31376e3cf
 size 140162704

build/torch210-cxx11-cu126-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "paged-attention",
-  "id": "_paged_attention_cuda_86f75d9",
+  "id": "_paged_attention_cuda_2567cd6",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

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

build/torch210-cxx11-cu128-x86_64-linux/{_paged_attention_cuda_86f75d9.abi3.so → _paged_attention_cuda_2567cd6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f827940d8899079daaa40f441d0b03198ca0196b28a66b888f2ea50fbf078c0a
+oid sha256:567e0ca0a3bd4f76e932894435215b92764fdb4e481463e81801c8a259db87dd
 size 167726096

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

@@ -1,6 +1,6 @@
 {
   "name": "paged-attention",
-  "id": "_paged_attention_cuda_86f75d9",
+  "id": "_paged_attention_cuda_2567cd6",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

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

build/torch210-cxx11-cu130-x86_64-linux/{_paged_attention_cuda_86f75d9.abi3.so → _paged_attention_cuda_2567cd6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:267cdf92780bf9c31872fb9e91a87f5ad16dc7978e44233cd83cf22c76d2c313
+oid sha256:f7b62d68a6276199ed7e179425c7bf434f3198b1cfa73d947cd1466f562edd99
 size 86560024

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

@@ -1,6 +1,6 @@
 {
   "name": "paged-attention",
-  "id": "_paged_attention_cuda_86f75d9",
+  "id": "_paged_attention_cuda_2567cd6",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch210-cxx11-rocm70-x86_64-linux/_ops.py

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

build/{torch210-cxx11-rocm71-x86_64-linux/_paged_attention_rocm_86f75d9.abi3.so → torch210-cxx11-rocm70-x86_64-linux/_paged_attention_rocm_2567cd6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7252bc2aec0ffa6e9fd79fa24c7f2ec4556538e84ac099467cdd91d97ae7045e
-size 58971840
+oid sha256:84ffe2ab8eaf52cfb8b656f0e4a7405f818d2d3837986862ba75c6efcf89fe7b
+size 84940576

build/torch210-cxx11-rocm70-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "paged-attention",
-  "id": "_paged_attention_rocm_86f75d9",
+  "id": "_paged_attention_rocm_2567cd6",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -10,10 +10,13 @@
       "gfx1030",
       "gfx1100",
       "gfx1101",
+      "gfx1200",
+      "gfx1201",
       "gfx906",
       "gfx908",
       "gfx90a",
-      "gfx942"
+      "gfx942",
+      "gfx950"
     ]
   }
 }

build/torch210-cxx11-rocm71-x86_64-linux/_ops.py

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

build/{torch210-cxx11-rocm70-x86_64-linux/_paged_attention_rocm_86f75d9.abi3.so → torch210-cxx11-rocm71-x86_64-linux/_paged_attention_rocm_2567cd6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c9198dc28f1a1f1c9b82c462bb2bc60aeb3e09b43295ed36aefc66d7e7af74f2
-size 58992416
+oid sha256:86fc9985afc44618710a543609c0694cf30bf2812c8aa0d0d5fb3048315d4f78
+size 84911808

build/torch210-cxx11-rocm71-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "paged-attention",
-  "id": "_paged_attention_rocm_86f75d9",
+  "id": "_paged_attention_rocm_2567cd6",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -10,10 +10,13 @@
       "gfx1030",
       "gfx1100",
       "gfx1101",
+      "gfx1200",
+      "gfx1201",
       "gfx906",
       "gfx908",
       "gfx90a",
-      "gfx942"
+      "gfx942",
+      "gfx950"
     ]
   }
 }

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

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

build/torch211-cxx11-cu126-x86_64-linux/{_paged_attention_cuda_86f75d9.abi3.so → _paged_attention_cuda_2567cd6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6f93258416e403f339435e72ae67787fb880962f5759dfaec3ab9202ed932b37
+oid sha256:9f245ffbd19fd92706e7332f1bdc35da933080407fc085e8b0fdd47007e1ee29
 size 140147224

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

@@ -1,6 +1,6 @@
 {
   "name": "paged-attention",
-  "id": "_paged_attention_cuda_86f75d9",
+  "id": "_paged_attention_cuda_2567cd6",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

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

build/torch211-cxx11-cu128-x86_64-linux/{_paged_attention_cuda_86f75d9.abi3.so → _paged_attention_cuda_2567cd6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5978826c7beb813daca2a0a9d596cd254c24c181b2c3922ed3b17771357bc20e
+oid sha256:45e8db1a82efeed8114df045332d9cf9cc5897896b46997eb834b926b8457529
 size 167710616

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

@@ -1,6 +1,6 @@
 {
   "name": "paged-attention",
-  "id": "_paged_attention_cuda_86f75d9",
+  "id": "_paged_attention_cuda_2567cd6",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

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

build/torch211-cxx11-cu130-x86_64-linux/{_paged_attention_cuda_86f75d9.abi3.so → _paged_attention_cuda_2567cd6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b897c1185c95f89f2dcc3a6fc13ea9a486d5939fc03574930689b4b89037efeb
+oid sha256:e97cfe7c0fb322c1ddda6253cfa6462f7320b82db874a93aee161857dfa3ee28
 size 86548640

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

@@ -1,6 +1,6 @@
 {
   "name": "paged-attention",
-  "id": "_paged_attention_cuda_86f75d9",
+  "id": "_paged_attention_cuda_2567cd6",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

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

build/torch211-cxx11-rocm71-x86_64-linux/_paged_attention_rocm_2567cd6.abi3.so

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

build/torch211-cxx11-rocm71-x86_64-linux/_paged_attention_rocm_86f75d9.abi3.so

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

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

@@ -1,6 +1,6 @@
 {
   "name": "paged-attention",
-  "id": "_paged_attention_rocm_86f75d9",
+  "id": "_paged_attention_rocm_2567cd6",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -10,10 +10,13 @@
       "gfx1030",
       "gfx1100",
       "gfx1101",
+      "gfx1200",
+      "gfx1201",
       "gfx906",
       "gfx908",
       "gfx90a",
-      "gfx942"
+      "gfx942",
+      "gfx950"
     ]
   }
 }

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

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

build/torch211-cxx11-rocm72-x86_64-linux/_paged_attention_rocm_2567cd6.abi3.so

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

build/torch211-cxx11-rocm72-x86_64-linux/_paged_attention_rocm_86f75d9.abi3.so

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

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

@@ -1,6 +1,6 @@
 {
   "name": "paged-attention",
-  "id": "_paged_attention_rocm_86f75d9",
+  "id": "_paged_attention_rocm_2567cd6",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -10,10 +10,13 @@
       "gfx1030",
       "gfx1100",
       "gfx1101",
+      "gfx1200",
+      "gfx1201",
       "gfx906",
       "gfx908",
       "gfx90a",
-      "gfx942"
+      "gfx942",
+      "gfx950"
     ]
   }
 }

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

@@ -0,0 +1,21 @@
+from ._custom_ops import (
+    convert_fp8,
+    copy_blocks,
+    paged_attention_v1,
+    paged_attention_v2,
+    reshape_and_cache,
+    reshape_and_cache_flash,
+    swap_blocks,
+)
+from ._ops import ops
+
+__all__ = [
+    "convert_fp8",
+    "copy_blocks",
+    "ops",
+    "paged_attention_v1",
+    "paged_attention_v2",
+    "reshape_and_cache",
+    "reshape_and_cache_flash",
+    "swap_blocks",
+]

build/torch212-cxx11-cu126-x86_64-linux/_custom_ops.py

@@ -0,0 +1,173 @@
+from typing import List, Optional
+
+import torch
+
+from ._ops import ops
+
+
+# page attention ops
+def paged_attention_v1(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    num_kv_heads: int,
+    scale: float,
+    block_tables: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    max_seq_len: int,
+    alibi_slopes: Optional[torch.Tensor],
+    kv_cache_dtype: str,
+    k_scale: float,
+    v_scale: float,
+    tp_rank: int = 0,
+    blocksparse_local_blocks: int = 0,
+    blocksparse_vert_stride: int = 0,
+    blocksparse_block_size: int = 64,
+    blocksparse_head_sliding_step: int = 0,
+) -> None:
+    ops.paged_attention_v1(
+        out,
+        query,
+        key_cache,
+        value_cache,
+        num_kv_heads,
+        scale,
+        block_tables,
+        seq_lens,
+        block_size,
+        max_seq_len,
+        alibi_slopes,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+        tp_rank,
+        blocksparse_local_blocks,
+        blocksparse_vert_stride,
+        blocksparse_block_size,
+        blocksparse_head_sliding_step,
+    )
+
+
+def paged_attention_v2(
+    out: torch.Tensor,
+    exp_sum: torch.Tensor,
+    max_logits: torch.Tensor,
+    tmp_out: torch.Tensor,
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    num_kv_heads: int,
+    scale: float,
+    block_tables: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    max_seq_len: int,
+    alibi_slopes: Optional[torch.Tensor],
+    kv_cache_dtype: str,
+    k_scale: float,
+    v_scale: float,
+    tp_rank: int = 0,
+    blocksparse_local_blocks: int = 0,
+    blocksparse_vert_stride: int = 0,
+    blocksparse_block_size: int = 64,
+    blocksparse_head_sliding_step: int = 0,
+) -> None:
+    ops.paged_attention_v2(
+        out,
+        exp_sum,
+        max_logits,
+        tmp_out,
+        query,
+        key_cache,
+        value_cache,
+        num_kv_heads,
+        scale,
+        block_tables,
+        seq_lens,
+        block_size,
+        max_seq_len,
+        alibi_slopes,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+        tp_rank,
+        blocksparse_local_blocks,
+        blocksparse_vert_stride,
+        blocksparse_block_size,
+        blocksparse_head_sliding_step,
+    )
+
+
+def reshape_and_cache(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    kv_cache_dtype: str,
+    k_scale: float,
+    v_scale: float,
+) -> None:
+    ops.reshape_and_cache(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+    )
+
+
+def reshape_and_cache_flash(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    kv_cache_dtype: str,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+) -> None:
+    ops.reshape_and_cache_flash(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+    )
+
+
+def copy_blocks(
+    key_caches: List[torch.Tensor],
+    value_caches: List[torch.Tensor],
+    block_mapping: torch.Tensor,
+) -> None:
+    ops.copy_blocks(key_caches, value_caches, block_mapping)
+
+
+def swap_blocks(
+    src: torch.Tensor, dst: torch.Tensor, block_mapping: torch.Tensor
+) -> None:
+    ops.swap_blocks(src, dst, block_mapping)
+
+
+def convert_fp8(
+    output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8"
+) -> None:
+    ops.convert_fp8(output, input, scale, kv_dtype)
+
+
+__all__ = [
+    "convert_fp8",
+    "paged_attention_v1",
+    "paged_attention_v2",
+    "reshape_and_cache",
+    "copy_blocks",
+]

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

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

build/torch212-cxx11-cu126-x86_64-linux/_paged_attention_cuda_2567cd6.abi3.so

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

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

@@ -0,0 +1,20 @@
+{
+  "name": "paged-attention",
+  "id": "_paged_attention_cuda_2567cd6",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "7.0",
+      "7.2",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0+PTX"
+    ]
+  }
+}

build/torch212-cxx11-cu126-x86_64-linux/paged_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-cu126-x86_64-linux/platforms.py

@@ -0,0 +1,92 @@
+import os
+import random
+from abc import ABC, abstractmethod
+from functools import lru_cache, wraps
+from typing import Callable, ParamSpec, TypeVar
+
+import numpy as np
+import torch
+
+IS_ROCM = torch.version.hip is not None
+IS_MPS = torch.backends.mps.is_available()
+
+
+class Platform(ABC):
+    @classmethod
+    def seed_everything(cls, seed: int) -> None:
+        """
+        Set the seed of each random module.
+        `torch.manual_seed` will set seed on all devices.
+
+        Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20
+        """
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+
+    @abstractmethod
+    def get_device_name(self, device_id: int = 0) -> str: ...
+
+    @abstractmethod
+    def is_cuda(self) -> bool: ...
+
+    @abstractmethod
+    def is_rocm(self) -> bool: ...
+
+    @abstractmethod
+    def is_mps(self) -> bool: ...
+
+
+class CudaPlatform(Platform):
+    @classmethod
+    @lru_cache(maxsize=8)
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return torch.cuda.get_device_name(0)
+
+    def is_cuda(self) -> bool:
+        return True
+
+    def is_rocm(self) -> bool:
+        return False
+
+    def is_mps(self) -> bool:
+        return False
+
+
+class RocmPlatform(Platform):
+    @classmethod
+    @lru_cache(maxsize=8)
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return torch.cuda.get_device_name(device_id)
+
+    def is_cuda(self) -> bool:
+        return False
+
+    def is_rocm(self) -> bool:
+        return True
+
+    def is_mps(self) -> bool:
+        return False
+
+
+class MpsPlatform(Platform):
+    @classmethod
+    @lru_cache(maxsize=8)
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return torch.cuda.get_device_name(device_id)
+
+    def is_cuda(self) -> bool:
+        return False
+
+    def is_rocm(self) -> bool:
+        return False
+
+    def is_mps(self) -> bool:
+        return True
+
+current_platform = (
+    RocmPlatform() if IS_ROCM else
+    MpsPlatform() if IS_MPS else
+    CudaPlatform() if torch.cuda.is_available() else
+    None
+)

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

@@ -0,0 +1,21 @@
+from ._custom_ops import (
+    convert_fp8,
+    copy_blocks,
+    paged_attention_v1,
+    paged_attention_v2,
+    reshape_and_cache,
+    reshape_and_cache_flash,
+    swap_blocks,
+)
+from ._ops import ops
+
+__all__ = [
+    "convert_fp8",
+    "copy_blocks",
+    "ops",
+    "paged_attention_v1",
+    "paged_attention_v2",
+    "reshape_and_cache",
+    "reshape_and_cache_flash",
+    "swap_blocks",
+]

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

@@ -0,0 +1,173 @@
+from typing import List, Optional
+
+import torch
+
+from ._ops import ops
+
+
+# page attention ops
+def paged_attention_v1(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    num_kv_heads: int,
+    scale: float,
+    block_tables: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    max_seq_len: int,
+    alibi_slopes: Optional[torch.Tensor],
+    kv_cache_dtype: str,
+    k_scale: float,
+    v_scale: float,
+    tp_rank: int = 0,
+    blocksparse_local_blocks: int = 0,
+    blocksparse_vert_stride: int = 0,
+    blocksparse_block_size: int = 64,
+    blocksparse_head_sliding_step: int = 0,
+) -> None:
+    ops.paged_attention_v1(
+        out,
+        query,
+        key_cache,
+        value_cache,
+        num_kv_heads,
+        scale,
+        block_tables,
+        seq_lens,
+        block_size,
+        max_seq_len,
+        alibi_slopes,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+        tp_rank,
+        blocksparse_local_blocks,
+        blocksparse_vert_stride,
+        blocksparse_block_size,
+        blocksparse_head_sliding_step,
+    )
+
+
+def paged_attention_v2(
+    out: torch.Tensor,
+    exp_sum: torch.Tensor,
+    max_logits: torch.Tensor,
+    tmp_out: torch.Tensor,
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    num_kv_heads: int,
+    scale: float,
+    block_tables: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    max_seq_len: int,
+    alibi_slopes: Optional[torch.Tensor],
+    kv_cache_dtype: str,
+    k_scale: float,
+    v_scale: float,
+    tp_rank: int = 0,
+    blocksparse_local_blocks: int = 0,
+    blocksparse_vert_stride: int = 0,
+    blocksparse_block_size: int = 64,
+    blocksparse_head_sliding_step: int = 0,
+) -> None:
+    ops.paged_attention_v2(
+        out,
+        exp_sum,
+        max_logits,
+        tmp_out,
+        query,
+        key_cache,
+        value_cache,
+        num_kv_heads,
+        scale,
+        block_tables,
+        seq_lens,
+        block_size,
+        max_seq_len,
+        alibi_slopes,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+        tp_rank,
+        blocksparse_local_blocks,
+        blocksparse_vert_stride,
+        blocksparse_block_size,
+        blocksparse_head_sliding_step,
+    )
+
+
+def reshape_and_cache(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    kv_cache_dtype: str,
+    k_scale: float,
+    v_scale: float,
+) -> None:
+    ops.reshape_and_cache(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+    )
+
+
+def reshape_and_cache_flash(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    kv_cache_dtype: str,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+) -> None:
+    ops.reshape_and_cache_flash(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+    )
+
+
+def copy_blocks(
+    key_caches: List[torch.Tensor],
+    value_caches: List[torch.Tensor],
+    block_mapping: torch.Tensor,
+) -> None:
+    ops.copy_blocks(key_caches, value_caches, block_mapping)
+
+
+def swap_blocks(
+    src: torch.Tensor, dst: torch.Tensor, block_mapping: torch.Tensor
+) -> None:
+    ops.swap_blocks(src, dst, block_mapping)
+
+
+def convert_fp8(
+    output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8"
+) -> None:
+    ops.convert_fp8(output, input, scale, kv_dtype)
+
+
+__all__ = [
+    "convert_fp8",
+    "paged_attention_v1",
+    "paged_attention_v2",
+    "reshape_and_cache",
+    "copy_blocks",
+]

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

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

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

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

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

@@ -0,0 +1,21 @@
+{
+  "name": "paged-attention",
+  "id": "_paged_attention_cuda_2567cd6",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch212-cxx11-cu130-x86_64-linux/paged_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/platforms.py

@@ -0,0 +1,92 @@
+import os
+import random
+from abc import ABC, abstractmethod
+from functools import lru_cache, wraps
+from typing import Callable, ParamSpec, TypeVar
+
+import numpy as np
+import torch
+
+IS_ROCM = torch.version.hip is not None
+IS_MPS = torch.backends.mps.is_available()
+
+
+class Platform(ABC):
+    @classmethod
+    def seed_everything(cls, seed: int) -> None:
+        """
+        Set the seed of each random module.
+        `torch.manual_seed` will set seed on all devices.
+
+        Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20
+        """
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+
+    @abstractmethod
+    def get_device_name(self, device_id: int = 0) -> str: ...
+
+    @abstractmethod
+    def is_cuda(self) -> bool: ...
+
+    @abstractmethod
+    def is_rocm(self) -> bool: ...
+
+    @abstractmethod
+    def is_mps(self) -> bool: ...
+
+
+class CudaPlatform(Platform):
+    @classmethod
+    @lru_cache(maxsize=8)
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return torch.cuda.get_device_name(0)
+
+    def is_cuda(self) -> bool:
+        return True
+
+    def is_rocm(self) -> bool:
+        return False
+
+    def is_mps(self) -> bool:
+        return False
+
+
+class RocmPlatform(Platform):
+    @classmethod
+    @lru_cache(maxsize=8)
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return torch.cuda.get_device_name(device_id)
+
+    def is_cuda(self) -> bool:
+        return False
+
+    def is_rocm(self) -> bool:
+        return True
+
+    def is_mps(self) -> bool:
+        return False
+
+
+class MpsPlatform(Platform):
+    @classmethod
+    @lru_cache(maxsize=8)
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return torch.cuda.get_device_name(device_id)
+
+    def is_cuda(self) -> bool:
+        return False
+
+    def is_rocm(self) -> bool:
+        return False
+
+    def is_mps(self) -> bool:
+        return True
+
+current_platform = (
+    RocmPlatform() if IS_ROCM else
+    MpsPlatform() if IS_MPS else
+    CudaPlatform() if torch.cuda.is_available() else
+    None
+)

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

@@ -0,0 +1,21 @@
+from ._custom_ops import (
+    convert_fp8,
+    copy_blocks,
+    paged_attention_v1,
+    paged_attention_v2,
+    reshape_and_cache,
+    reshape_and_cache_flash,
+    swap_blocks,
+)
+from ._ops import ops
+
+__all__ = [
+    "convert_fp8",
+    "copy_blocks",
+    "ops",
+    "paged_attention_v1",
+    "paged_attention_v2",
+    "reshape_and_cache",
+    "reshape_and_cache_flash",
+    "swap_blocks",
+]

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

@@ -0,0 +1,173 @@
+from typing import List, Optional
+
+import torch
+
+from ._ops import ops
+
+
+# page attention ops
+def paged_attention_v1(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    num_kv_heads: int,
+    scale: float,
+    block_tables: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    max_seq_len: int,
+    alibi_slopes: Optional[torch.Tensor],
+    kv_cache_dtype: str,
+    k_scale: float,
+    v_scale: float,
+    tp_rank: int = 0,
+    blocksparse_local_blocks: int = 0,
+    blocksparse_vert_stride: int = 0,
+    blocksparse_block_size: int = 64,
+    blocksparse_head_sliding_step: int = 0,
+) -> None:
+    ops.paged_attention_v1(
+        out,
+        query,
+        key_cache,
+        value_cache,
+        num_kv_heads,
+        scale,
+        block_tables,
+        seq_lens,
+        block_size,
+        max_seq_len,
+        alibi_slopes,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+        tp_rank,
+        blocksparse_local_blocks,
+        blocksparse_vert_stride,
+        blocksparse_block_size,
+        blocksparse_head_sliding_step,
+    )
+
+
+def paged_attention_v2(
+    out: torch.Tensor,
+    exp_sum: torch.Tensor,
+    max_logits: torch.Tensor,
+    tmp_out: torch.Tensor,
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    num_kv_heads: int,
+    scale: float,
+    block_tables: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    max_seq_len: int,
+    alibi_slopes: Optional[torch.Tensor],
+    kv_cache_dtype: str,
+    k_scale: float,
+    v_scale: float,
+    tp_rank: int = 0,
+    blocksparse_local_blocks: int = 0,
+    blocksparse_vert_stride: int = 0,
+    blocksparse_block_size: int = 64,
+    blocksparse_head_sliding_step: int = 0,
+) -> None:
+    ops.paged_attention_v2(
+        out,
+        exp_sum,
+        max_logits,
+        tmp_out,
+        query,
+        key_cache,
+        value_cache,
+        num_kv_heads,
+        scale,
+        block_tables,
+        seq_lens,
+        block_size,
+        max_seq_len,
+        alibi_slopes,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+        tp_rank,
+        blocksparse_local_blocks,
+        blocksparse_vert_stride,
+        blocksparse_block_size,
+        blocksparse_head_sliding_step,
+    )
+
+
+def reshape_and_cache(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    kv_cache_dtype: str,
+    k_scale: float,
+    v_scale: float,
+) -> None:
+    ops.reshape_and_cache(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+    )
+
+
+def reshape_and_cache_flash(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    kv_cache_dtype: str,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+) -> None:
+    ops.reshape_and_cache_flash(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        slot_mapping,
+        kv_cache_dtype,
+        k_scale,
+        v_scale,
+    )
+
+
+def copy_blocks(
+    key_caches: List[torch.Tensor],
+    value_caches: List[torch.Tensor],
+    block_mapping: torch.Tensor,
+) -> None:
+    ops.copy_blocks(key_caches, value_caches, block_mapping)
+
+
+def swap_blocks(
+    src: torch.Tensor, dst: torch.Tensor, block_mapping: torch.Tensor
+) -> None:
+    ops.swap_blocks(src, dst, block_mapping)
+
+
+def convert_fp8(
+    output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8"
+) -> None:
+    ops.convert_fp8(output, input, scale, kv_dtype)
+
+
+__all__ = [
+    "convert_fp8",
+    "paged_attention_v1",
+    "paged_attention_v2",
+    "reshape_and_cache",
+    "copy_blocks",
+]

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

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

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

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