Build

build/torch25-cxx11-cu118-x86_64-linux/paged_attention/_ops.py

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

build/torch25-cxx11-cu118-x86_64-linux/paged_attention/{_paged_attention_fao6f4gjjrpl6.abi3.so → _paged_attention_3f98f45.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dd3f60d12b061ef8cc4adb9bd2d60539bc233075c007af99b995e4c12bb5df98
-size 78886416
+oid sha256:f6d91b11a795c99585e3223373c80aebadb83c704d1cfe91d116b36c90db5177
+size 74745336

build/torch25-cxx11-cu121-x86_64-linux/paged_attention/_ops.py

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

build/torch25-cxx11-cu121-x86_64-linux/paged_attention/{_paged_attention_eo7ts45r6k64y.abi3.so → _paged_attention_3f98f45.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c4f9be8ea04e306483bd77ee84a70be4b1e2bbe1e942766f398cb41afe8cb3f3
-size 75398840
+oid sha256:7f771fc8f788563819d87b06da86c7a6169f923ee62a69f2ef62331a8d11fee5
+size 75603608

build/torch25-cxx11-cu124-x86_64-linux/paged_attention/_ops.py

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

build/torch25-cxx11-cu124-x86_64-linux/paged_attention/{_paged_attention_5odgyxqhwqtv2.abi3.so → _paged_attention_3f98f45.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4fbf3b6cf357130fb9bb6e68cab38cbcbf485147d9dcce1a9a7d7c02fb2e1af8
-size 75568352
+oid sha256:33d0933dd8198257c015f823eecaec5e2f4c00d2debb0ecd8de2724dbb5a9958
+size 75715776

build/torch25-cxx98-cu118-x86_64-linux/paged_attention/_ops.py

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

build/torch25-cxx98-cu118-x86_64-linux/paged_attention/{_paged_attention_uy2moinaww2jc.abi3.so → _paged_attention_3f98f45.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:546d3cb8850575fdd3dd198614d7c64894517dfc124ecd0b384f1918940e9db0
-size 78850640
+oid sha256:708d455ea0861e805df9291b1fe30094d8dfb43470b599791ab59a33c52b30b8
+size 74709552

build/torch25-cxx98-cu121-x86_64-linux/paged_attention/_ops.py

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

build/torch25-cxx98-cu121-x86_64-linux/paged_attention/_paged_attention_35dt23tewn2p2.abi3.so

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

build/torch25-cxx98-cu121-x86_64-linux/paged_attention/_paged_attention_3f98f45.abi3.so

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

build/torch25-cxx98-cu124-x86_64-linux/paged_attention/_ops.py

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

build/torch25-cxx98-cu124-x86_64-linux/paged_attention/_paged_attention_3f98f45.abi3.so

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

build/torch25-cxx98-cu124-x86_64-linux/paged_attention/_paged_attention_fhq57q56w3m5o.abi3.so

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

build/torch26-cxx11-cu118-x86_64-linux/paged_attention/_ops.py

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

build/torch26-cxx11-cu118-x86_64-linux/paged_attention/_paged_attention_3f98f45.abi3.so

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

build/torch26-cxx11-cu118-x86_64-linux/paged_attention/_paged_attention_xvepb4loq5mm2.abi3.so

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

build/torch26-cxx11-cu124-x86_64-linux/paged_attention/_ops.py

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

build/torch26-cxx11-cu124-x86_64-linux/paged_attention/_paged_attention_3f98f45.abi3.so

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

build/torch26-cxx11-cu124-x86_64-linux/paged_attention/_paged_attention_uyfdujhnc2xoe.abi3.so

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

build/torch26-cxx11-cu126-x86_64-linux/paged_attention/_ops.py

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

build/torch26-cxx11-cu126-x86_64-linux/paged_attention/_paged_attention_3f98f45.abi3.so

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

build/torch26-cxx11-cu126-x86_64-linux/paged_attention/_paged_attention_pervvqmod6pi4.abi3.so

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

build/torch26-cxx98-cu118-x86_64-linux/paged_attention/_ops.py

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

build/torch26-cxx98-cu118-x86_64-linux/paged_attention/_paged_attention_24rowhxd5ebcc.abi3.so

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

build/torch26-cxx98-cu118-x86_64-linux/paged_attention/_paged_attention_3f98f45.abi3.so

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

build/torch26-cxx98-cu124-x86_64-linux/paged_attention/_ops.py

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

build/torch26-cxx98-cu124-x86_64-linux/paged_attention/_paged_attention_3f98f45.abi3.so

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

build/torch26-cxx98-cu124-x86_64-linux/paged_attention/_paged_attention_5yleoqr3zje4w.abi3.so

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

build/torch26-cxx98-cu126-x86_64-linux/paged_attention/_ops.py

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

build/torch26-cxx98-cu126-x86_64-linux/paged_attention/_paged_attention_3f98f45.abi3.so

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

build/torch26-cxx98-cu126-x86_64-linux/paged_attention/_paged_attention_3rbp7xipfucgo.abi3.so

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

build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__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/torch27-cxx11-cu118-x86_64-linux/paged_attention/_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/torch27-cxx11-cu118-x86_64-linux/paged_attention/_ops.py

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

build/torch27-cxx11-cu118-x86_64-linux/paged_attention/_paged_attention_3f98f45.abi3.so

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

build/torch27-cxx11-cu118-x86_64-linux/paged_attention/platforms.py

@@ -0,0 +1,62 @@
+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
+
+
+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: ...
+
+
+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
+
+
+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
+
+
+current_platform = RocmPlatform() if IS_ROCM else CudaPlatform()

build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__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/torch27-cxx11-cu126-x86_64-linux/paged_attention/_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/torch27-cxx11-cu126-x86_64-linux/paged_attention/_ops.py

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

build/torch27-cxx11-cu126-x86_64-linux/paged_attention/_paged_attention_3f98f45.abi3.so

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

build/torch27-cxx11-cu126-x86_64-linux/paged_attention/platforms.py

@@ -0,0 +1,62 @@
+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
+
+
+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: ...
+
+
+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
+
+
+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
+
+
+current_platform = RocmPlatform() if IS_ROCM else CudaPlatform()

build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__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/torch27-cxx11-cu128-x86_64-linux/paged_attention/_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/torch27-cxx11-cu128-x86_64-linux/paged_attention/_ops.py

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

build/torch27-cxx11-cu128-x86_64-linux/paged_attention/_paged_attention_3f98f45.abi3.so

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

build/torch27-cxx11-cu128-x86_64-linux/paged_attention/platforms.py

@@ -0,0 +1,62 @@
+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
+
+
+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: ...
+
+
+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
+
+
+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
+
+
+current_platform = RocmPlatform() if IS_ROCM else CudaPlatform()