diff --git a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__init__.py b/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index 3a42f6c23989d05e0c4180a8e07ce104a7ecd99c..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc b/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc deleted file mode 100644 index 9c698ce4a5f2b3db77ed9861efc143265073f6a8..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index f4cbb5b24a571485fb2d7712b46631a1eedf0b41..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/_custom_ops.py b/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/_ops.py b/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/_ops.py deleted file mode 100644 index a883d5cb5d351eb18b40e1d4b621a7d1544385ab..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_0041e3f -ops = torch.ops._paged_attention_0041e3f - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_0041e3f::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so b/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so deleted file mode 100755 index 71626cdbf291f8b39d80a6b511439c2148f25112..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b899f376425b8d7a213b8d26909d84ffb1213e03d3e5b33675e9408426747501 -size 113844912 diff --git a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/platforms.py b/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/paged_attention/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__init__.py b/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index 428ec0cbce6fbd264193a66297408d47d9f9882b..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc b/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc deleted file mode 100644 index c4257fb3f872d1ec4e346d6f3c5ec2046ab82f36..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index 987598889f8379c07efa94bad6a3b6f17dec9860..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/_custom_ops.py b/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/_ops.py b/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/_ops.py deleted file mode 100644 index a883d5cb5d351eb18b40e1d4b621a7d1544385ab..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_0041e3f -ops = torch.ops._paged_attention_0041e3f - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_0041e3f::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so b/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so deleted file mode 100755 index d68462effa6cb3bc865b2cdca6137261d51c2edc..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f98a11d55d75ee841515253dfd16b5dd17a811925a445ef765b90b4b56a10e35 -size 110732296 diff --git a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/platforms.py b/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/paged_attention/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__init__.py b/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index 4426d89b40d2875c6242745d3a3404dc105dc340..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc b/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc deleted file mode 100644 index 76058206a88f8f66740efd7b611f125a5f3e1eee..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index 0056674c98c61fb9850920022abc1bbebebdb2e2..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/_custom_ops.py b/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/_ops.py b/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/_ops.py deleted file mode 100644 index a5e9f56ff02bdcf047fd97ef2a46b10d6fb5e5eb..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_b4c51e9 -ops = torch.ops._paged_attention_b4c51e9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_b4c51e9::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so b/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so deleted file mode 100755 index 31455c6c5a159469575f362d673d272a341e4b19..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9cba9e7b3cf9de722c5dc8f56a533542521fdef05a290a2b4db13948b0f1eca4 -size 138172880 diff --git a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/platforms.py b/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-aarch64-linux/paged_attention/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__init__.py b/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index 7c13b4dd0a5c62749c4fffaed83061e327465cfa..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc b/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc deleted file mode 100644 index fb10ad97588c556a857e62d03549f1c72efeced9..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index dbe695e89609fad8076312195af5a04690a744ee..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/_custom_ops.py b/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/_ops.py b/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/_ops.py deleted file mode 100644 index a883d5cb5d351eb18b40e1d4b621a7d1544385ab..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_0041e3f -ops = torch.ops._paged_attention_0041e3f - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_0041e3f::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so b/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so deleted file mode 100755 index e2fbd54b5c6e0379de056949d3d7991b31b546e6..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:15f5767adca82992a20bb3b0ff047e5cd076c386f1bdff70f63c3052c47e1ac3 -size 138291040 diff --git a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/platforms.py b/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/paged_attention/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__init__.py b/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc b/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index 8a487f6dd62db7a85db734773b9718a43478a689..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc b/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc deleted file mode 100644 index b8ae1d42b0f5f3fb681398076bef151d34cc9325..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc b/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index 7a4ea49b96acc139399eedbd529c4af03da4944b..0000000000000000000000000000000000000000 Binary files a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/_custom_ops.py b/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/_ops.py b/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/_ops.py deleted file mode 100644 index a883d5cb5d351eb18b40e1d4b621a7d1544385ab..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_0041e3f -ops = torch.ops._paged_attention_0041e3f - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_0041e3f::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so b/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so deleted file mode 100755 index f8326ce54482175c1763ff12b7789f71c61dc9ca..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/_paged_attention_0041e3f.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5955ede80b62f77cbcf76ffc3961b20af2d5a17ceb95c5c2fc06e7b12e9d3bec -size 120178304 diff --git a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/platforms.py b/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-rocm63-x86_64-linux/paged_attention/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch27-metal-aarch64-darwin/paged_attention/__init__.py b/build/torch27-metal-aarch64-darwin/paged_attention/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch27-metal-aarch64-darwin/paged_attention/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch27-metal-aarch64-darwin/paged_attention/_custom_ops.py b/build/torch27-metal-aarch64-darwin/paged_attention/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch27-metal-aarch64-darwin/paged_attention/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch27-metal-aarch64-darwin/paged_attention/_ops.py b/build/torch27-metal-aarch64-darwin/paged_attention/_ops.py deleted file mode 100644 index e6a82eda58b197517abf459ac6adaefaf817431b..0000000000000000000000000000000000000000 --- a/build/torch27-metal-aarch64-darwin/paged_attention/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_9678b89 -ops = torch.ops._paged_attention_9678b89 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_9678b89::{op_name}" \ No newline at end of file diff --git a/build/torch27-metal-aarch64-darwin/paged_attention/_paged_attention_9678b89.abi3.so b/build/torch27-metal-aarch64-darwin/paged_attention/_paged_attention_9678b89.abi3.so deleted file mode 100755 index d6c43e8f7b7b269de282330ba9c877243937a713..0000000000000000000000000000000000000000 --- a/build/torch27-metal-aarch64-darwin/paged_attention/_paged_attention_9678b89.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7a94cee9e553d2bdf8d47d0d9461c871b3e57a33cf6cb259807377f0d1b03c7d -size 214800 diff --git a/build/torch27-metal-aarch64-darwin/paged_attention/_paged_attention_9678b89.metallib b/build/torch27-metal-aarch64-darwin/paged_attention/_paged_attention_9678b89.metallib deleted file mode 100644 index 076fc55c592b25de05c25232c560b617986561cc..0000000000000000000000000000000000000000 --- a/build/torch27-metal-aarch64-darwin/paged_attention/_paged_attention_9678b89.metallib +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c46eaf21c96da70c5227b2566308a8ef73ae09abf303278f40070dd4326ba0be -size 4999876 diff --git a/build/torch27-metal-aarch64-darwin/paged_attention/platforms.py b/build/torch27-metal-aarch64-darwin/paged_attention/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch27-metal-aarch64-darwin/paged_attention/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch28-cxx11-cu126-x86_64-linux/__init__.py b/build/torch28-cxx11-cu126-x86_64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-cu126-x86_64-linux/_custom_ops.py b/build/torch28-cxx11-cu126-x86_64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-cu126-x86_64-linux/_ops.py b/build/torch28-cxx11-cu126-x86_64-linux/_ops.py deleted file mode 100644 index db595747ff801b0cf8021fcbb30693389cbebb94..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_7c06fa7 -ops = torch.ops._paged_attention_7c06fa7 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_7c06fa7::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu126-x86_64-linux/_paged_attention_7c06fa7.abi3.so b/build/torch28-cxx11-cu126-x86_64-linux/_paged_attention_7c06fa7.abi3.so deleted file mode 100644 index d846ed81c7f6d3218dee88b09cf9a283a304b8f6..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/_paged_attention_7c06fa7.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c3e8736d3b0efd5435d028d3f4429ebcc8fb95bf9dcaa3d7b191ebcea5ad20c2 -size 140143504 diff --git a/build/torch28-cxx11-cu126-x86_64-linux/metadata.json b/build/torch28-cxx11-cu126-x86_64-linux/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-cu126-x86_64-linux/paged_attention/__init__.py b/build/torch28-cxx11-cu126-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch28-cxx11-cu126-x86_64-linux/platforms.py b/build/torch28-cxx11-cu126-x86_64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu126-x86_64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch28-cxx11-cu128-x86_64-linux/__init__.py b/build/torch28-cxx11-cu128-x86_64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-cu128-x86_64-linux/_custom_ops.py b/build/torch28-cxx11-cu128-x86_64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-cu128-x86_64-linux/_ops.py b/build/torch28-cxx11-cu128-x86_64-linux/_ops.py deleted file mode 100644 index db595747ff801b0cf8021fcbb30693389cbebb94..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_7c06fa7 -ops = torch.ops._paged_attention_7c06fa7 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_7c06fa7::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu128-x86_64-linux/_paged_attention_7c06fa7.abi3.so b/build/torch28-cxx11-cu128-x86_64-linux/_paged_attention_7c06fa7.abi3.so deleted file mode 100644 index ab267a807aba88bf8eed5bc0adbec038e4b8c42e..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/_paged_attention_7c06fa7.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ba3a28b63cf7a420d5f1622beec6b8d916254484f8e8d1629f33ee2c5ae4a7a0 -size 167710768 diff --git a/build/torch28-cxx11-cu128-x86_64-linux/metadata.json b/build/torch28-cxx11-cu128-x86_64-linux/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-cu128-x86_64-linux/paged_attention/__init__.py b/build/torch28-cxx11-cu128-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch28-cxx11-cu128-x86_64-linux/platforms.py b/build/torch28-cxx11-cu128-x86_64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu128-x86_64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__init__.py b/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc b/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index f2029d1fa1af1cacf90cbe0e8a2a59b907940e8d..0000000000000000000000000000000000000000 Binary files a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc b/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc deleted file mode 100644 index 4f54c4683b6c5a5d05d435d40f546ac8cbf62e76..0000000000000000000000000000000000000000 Binary files a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc b/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index 095bc0f18dddbaa40b04468fa6cb9dde32b9285a..0000000000000000000000000000000000000000 Binary files a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/_custom_ops.py b/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/_ops.py b/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/_ops.py deleted file mode 100644 index a5e9f56ff02bdcf047fd97ef2a46b10d6fb5e5eb..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_b4c51e9 -ops = torch.ops._paged_attention_b4c51e9 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_b4c51e9::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so b/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so deleted file mode 100755 index f4f21766fa2c547675a1373926d281b4ec7a28fa..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:02238a0a4dacdbbf60eb9bb73a95332b448d5127f71149003c764f40595d9d06 -size 149841048 diff --git a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/platforms.py b/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-aarch64-linux/paged_attention/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch28-cxx11-cu129-x86_64-linux/__init__.py b/build/torch28-cxx11-cu129-x86_64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-cu129-x86_64-linux/_custom_ops.py b/build/torch28-cxx11-cu129-x86_64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-cu129-x86_64-linux/_ops.py b/build/torch28-cxx11-cu129-x86_64-linux/_ops.py deleted file mode 100644 index db595747ff801b0cf8021fcbb30693389cbebb94..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_7c06fa7 -ops = torch.ops._paged_attention_7c06fa7 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_7c06fa7::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-cu129-x86_64-linux/_paged_attention_7c06fa7.abi3.so b/build/torch28-cxx11-cu129-x86_64-linux/_paged_attention_7c06fa7.abi3.so deleted file mode 100644 index f7fbade89d1c107240d641411d712cf7829dbbe8..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/_paged_attention_7c06fa7.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b20a5cd5515e0a521c241bd4c7d659b82e6393217ebb86b7e8360d8b19949994 -size 182084600 diff --git a/build/torch28-cxx11-cu129-x86_64-linux/metadata.json b/build/torch28-cxx11-cu129-x86_64-linux/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-cu129-x86_64-linux/paged_attention/__init__.py b/build/torch28-cxx11-cu129-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch28-cxx11-cu129-x86_64-linux/platforms.py b/build/torch28-cxx11-cu129-x86_64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-cu129-x86_64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/__init__.py b/build/torch28-cxx11-rocm63-x86_64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/_custom_ops.py b/build/torch28-cxx11-rocm63-x86_64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/_ops.py b/build/torch28-cxx11-rocm63-x86_64-linux/_ops.py deleted file mode 100644 index db595747ff801b0cf8021fcbb30693389cbebb94..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_7c06fa7 -ops = torch.ops._paged_attention_7c06fa7 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_7c06fa7::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/_paged_attention_7c06fa7.abi3.so b/build/torch28-cxx11-rocm63-x86_64-linux/_paged_attention_7c06fa7.abi3.so deleted file mode 100644 index 0b4395723360f06cb1c5968a4c4818ab6085a2f5..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/_paged_attention_7c06fa7.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e61cf1b792a2479d48a2c886b30abf364efa435ced959dbfca03a619dbdd47e6 -size 120192848 diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/metadata.json b/build/torch28-cxx11-rocm63-x86_64-linux/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/paged_attention/__init__.py b/build/torch28-cxx11-rocm63-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch28-cxx11-rocm63-x86_64-linux/platforms.py b/build/torch28-cxx11-rocm63-x86_64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm63-x86_64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/__init__.py b/build/torch28-cxx11-rocm64-x86_64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/_custom_ops.py b/build/torch28-cxx11-rocm64-x86_64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/_ops.py b/build/torch28-cxx11-rocm64-x86_64-linux/_ops.py deleted file mode 100644 index db595747ff801b0cf8021fcbb30693389cbebb94..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_7c06fa7 -ops = torch.ops._paged_attention_7c06fa7 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_7c06fa7::{op_name}" \ No newline at end of file diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/_paged_attention_7c06fa7.abi3.so b/build/torch28-cxx11-rocm64-x86_64-linux/_paged_attention_7c06fa7.abi3.so deleted file mode 100644 index 8a198fa60c375cd9a63533504b065f74bd6667c9..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/_paged_attention_7c06fa7.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:39d897931a7f06c168722ff3ec836eb923d528f2648cf801e55473bc726a0e5b -size 121030544 diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/metadata.json b/build/torch28-cxx11-rocm64-x86_64-linux/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/paged_attention/__init__.py b/build/torch28-cxx11-rocm64-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch28-cxx11-rocm64-x86_64-linux/platforms.py b/build/torch28-cxx11-rocm64-x86_64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch28-cxx11-rocm64-x86_64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch28-metal-aarch64-darwin/__init__.py b/build/torch28-metal-aarch64-darwin/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch28-metal-aarch64-darwin/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch28-metal-aarch64-darwin/_custom_ops.py b/build/torch28-metal-aarch64-darwin/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch28-metal-aarch64-darwin/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch28-metal-aarch64-darwin/_ops.py b/build/torch28-metal-aarch64-darwin/_ops.py deleted file mode 100644 index db595747ff801b0cf8021fcbb30693389cbebb94..0000000000000000000000000000000000000000 --- a/build/torch28-metal-aarch64-darwin/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_7c06fa7 -ops = torch.ops._paged_attention_7c06fa7 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_7c06fa7::{op_name}" \ No newline at end of file diff --git a/build/torch28-metal-aarch64-darwin/_paged_attention_7c06fa7.abi3.so b/build/torch28-metal-aarch64-darwin/_paged_attention_7c06fa7.abi3.so deleted file mode 100644 index 410ba8a85bdcc005471f889cbb19cf47e9ab3270..0000000000000000000000000000000000000000 --- a/build/torch28-metal-aarch64-darwin/_paged_attention_7c06fa7.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0626edd9c58bef7c74b32cbdb73219d4e49fb650a0098726c4856dab33888306 -size 14892816 diff --git a/build/torch28-metal-aarch64-darwin/metadata.json b/build/torch28-metal-aarch64-darwin/metadata.json deleted file mode 100644 index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..0000000000000000000000000000000000000000 --- a/build/torch28-metal-aarch64-darwin/metadata.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "version": 1, - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch28-metal-aarch64-darwin/paged_attention/__init__.py b/build/torch28-metal-aarch64-darwin/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch28-metal-aarch64-darwin/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch28-metal-aarch64-darwin/platforms.py b/build/torch28-metal-aarch64-darwin/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch28-metal-aarch64-darwin/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/__init__.py b/build/torch29-cxx11-cu126-aarch64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu126-aarch64-linux/_custom_ops.py b/build/torch29-cxx11-cu126-aarch64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu126-aarch64-linux/_ops.py b/build/torch29-cxx11-cu126-aarch64-linux/_ops.py deleted file mode 100644 index 5415b95a504475517ce0b451729d8edaaf24042a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_cuda_c6404f1 -ops = torch.ops._paged_attention_cuda_c6404f1 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_cuda_c6404f1::{op_name}" diff --git a/build/torch29-cxx11-cu126-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so b/build/torch29-cxx11-cu126-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so deleted file mode 100644 index 1760b475fbbb0d8a67982aaccbd8d83fc598157c..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:beaa2bb81fd77f1a1bd29447e875044d741aa5069d88beff50234a45a0dee32e -size 140009952 diff --git a/build/torch29-cxx11-cu126-aarch64-linux/metadata.json b/build/torch29-cxx11-cu126-aarch64-linux/metadata.json deleted file mode 100644 index f5902b55ab0b2b561c0cf97567c9806c60839c7f..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/metadata.json +++ /dev/null @@ -1,18 +0,0 @@ -{ - "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" - ] - } -} diff --git a/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__init__.py b/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch29-cxx11-cu126-aarch64-linux/platforms.py b/build/torch29-cxx11-cu126-aarch64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-aarch64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch29-cxx11-cu126-x86_64-linux/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_custom_ops.py b/build/torch29-cxx11-cu126-x86_64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_ops.py b/build/torch29-cxx11-cu126-x86_64-linux/_ops.py deleted file mode 100644 index 5415b95a504475517ce0b451729d8edaaf24042a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_cuda_c6404f1 -ops = torch.ops._paged_attention_cuda_c6404f1 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_cuda_c6404f1::{op_name}" diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_paged_attention_cuda_c6404f1.abi3.so b/build/torch29-cxx11-cu126-x86_64-linux/_paged_attention_cuda_c6404f1.abi3.so deleted file mode 100644 index d5a2175146a85e7a9bfb143f4ad583834a0069fc..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/_paged_attention_cuda_c6404f1.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:89325f1b97efcf0d22e55ee4a68bb7590296190384a5af34702b355566b9140c -size 140143504 diff --git a/build/torch29-cxx11-cu126-x86_64-linux/metadata.json b/build/torch29-cxx11-cu126-x86_64-linux/metadata.json deleted file mode 100644 index f5902b55ab0b2b561c0cf97567c9806c60839c7f..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/metadata.json +++ /dev/null @@ -1,18 +0,0 @@ -{ - "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" - ] - } -} diff --git a/build/torch29-cxx11-cu126-x86_64-linux/paged_attention/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch29-cxx11-cu126-x86_64-linux/platforms.py b/build/torch29-cxx11-cu126-x86_64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu128-aarch64-linux/_custom_ops.py b/build/torch29-cxx11-cu128-aarch64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu128-aarch64-linux/_ops.py b/build/torch29-cxx11-cu128-aarch64-linux/_ops.py deleted file mode 100644 index 5415b95a504475517ce0b451729d8edaaf24042a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_cuda_c6404f1 -ops = torch.ops._paged_attention_cuda_c6404f1 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_cuda_c6404f1::{op_name}" diff --git a/build/torch29-cxx11-cu128-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so b/build/torch29-cxx11-cu128-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so deleted file mode 100644 index e4641642e57ed6726f75d4e867b337f0f01ad4f8..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d5708848d9e6113b481537edecd14c115844ab223acb66b64dcf78148cd7f2d1 -size 167599904 diff --git a/build/torch29-cxx11-cu128-aarch64-linux/metadata.json b/build/torch29-cxx11-cu128-aarch64-linux/metadata.json deleted file mode 100644 index 8b796af185fbbd8594fcd846949aa5fadc0ccdda..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/metadata.json +++ /dev/null @@ -1,21 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "10.1", - "12.0+PTX", - "7.0", - "7.2", - "7.5", - "8.0", - "8.6", - "8.7", - "8.9", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/platforms.py b/build/torch29-cxx11-cu128-aarch64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_custom_ops.py b/build/torch29-cxx11-cu128-x86_64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py b/build/torch29-cxx11-cu128-x86_64-linux/_ops.py deleted file mode 100644 index 5415b95a504475517ce0b451729d8edaaf24042a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_cuda_c6404f1 -ops = torch.ops._paged_attention_cuda_c6404f1 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_cuda_c6404f1::{op_name}" diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_paged_attention_cuda_c6404f1.abi3.so b/build/torch29-cxx11-cu128-x86_64-linux/_paged_attention_cuda_c6404f1.abi3.so deleted file mode 100644 index 8357452f06bd5c4993903fffd65accaaaadc78cf..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/_paged_attention_cuda_c6404f1.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fb574b6d123b5a419063568d05231969f4a70a5fb2d4903c9d1df60c8d5fc9a2 -size 167710768 diff --git a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json b/build/torch29-cxx11-cu128-x86_64-linux/metadata.json deleted file mode 100644 index 8b796af185fbbd8594fcd846949aa5fadc0ccdda..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json +++ /dev/null @@ -1,21 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "10.1", - "12.0+PTX", - "7.0", - "7.2", - "7.5", - "8.0", - "8.6", - "8.7", - "8.9", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu128-x86_64-linux/paged_attention/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/platforms.py b/build/torch29-cxx11-cu128-x86_64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/__init__.py b/build/torch29-cxx11-cu129-aarch64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu129-aarch64-linux/_custom_ops.py b/build/torch29-cxx11-cu129-aarch64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu129-aarch64-linux/_ops.py b/build/torch29-cxx11-cu129-aarch64-linux/_ops.py deleted file mode 100644 index 3317750fa0f3070e4bc61f664924506bcc137445..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_cuda_04f8f6e -ops = torch.ops._paged_attention_cuda_04f8f6e - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_cuda_04f8f6e::{op_name}" diff --git a/build/torch29-cxx11-cu129-aarch64-linux/_paged_attention_cuda_04f8f6e.abi3.so b/build/torch29-cxx11-cu129-aarch64-linux/_paged_attention_cuda_04f8f6e.abi3.so deleted file mode 100644 index d26cecdf7faa4aa2ce3e8405359da5916ba7aa1b..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/_paged_attention_cuda_04f8f6e.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a2281e8180dc6de2ef41a1ece7940d0b5b5af29b6a947777c43180c1fda1d99f -size 181953688 diff --git a/build/torch29-cxx11-cu129-aarch64-linux/metadata.json b/build/torch29-cxx11-cu129-aarch64-linux/metadata.json deleted file mode 100644 index 8b796af185fbbd8594fcd846949aa5fadc0ccdda..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/metadata.json +++ /dev/null @@ -1,21 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "10.1", - "12.0+PTX", - "7.0", - "7.2", - "7.5", - "8.0", - "8.6", - "8.7", - "8.9", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu129-aarch64-linux/paged_attention/__init__.py b/build/torch29-cxx11-cu129-aarch64-linux/paged_attention/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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"))) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/platforms.py b/build/torch29-cxx11-cu129-aarch64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu129-x86_64-linux/_custom_ops.py b/build/torch29-cxx11-cu129-x86_64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu129-x86_64-linux/_ops.py b/build/torch29-cxx11-cu129-x86_64-linux/_ops.py deleted file mode 100644 index 3317750fa0f3070e4bc61f664924506bcc137445..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_cuda_04f8f6e -ops = torch.ops._paged_attention_cuda_04f8f6e - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_cuda_04f8f6e::{op_name}" diff --git a/build/torch29-cxx11-cu129-x86_64-linux/_paged_attention_cuda_04f8f6e.abi3.so b/build/torch29-cxx11-cu129-x86_64-linux/_paged_attention_cuda_04f8f6e.abi3.so deleted file mode 100644 index e8f5ab28f8ae2fd99a18d720ad1d7c0e45f81e7e..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/_paged_attention_cuda_04f8f6e.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e35378b47b6101c82d9d227baf8402a8c42d46a8be6072ac20b52f2e528eaafb -size 182088696 diff --git a/build/torch29-cxx11-cu129-x86_64-linux/metadata.json b/build/torch29-cxx11-cu129-x86_64-linux/metadata.json deleted file mode 100644 index 8b796af185fbbd8594fcd846949aa5fadc0ccdda..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/metadata.json +++ /dev/null @@ -1,21 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0", - "10.1", - "12.0+PTX", - "7.0", - "7.2", - "7.5", - "8.0", - "8.6", - "8.7", - "8.9", - "9.0" - ] - } -} diff --git a/build/torch29-cxx11-cu129-x86_64-linux/paged_attention/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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"))) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/platforms.py b/build/torch29-cxx11-cu129-x86_64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu130-aarch64-linux/_custom_ops.py b/build/torch29-cxx11-cu130-aarch64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu130-aarch64-linux/_ops.py b/build/torch29-cxx11-cu130-aarch64-linux/_ops.py deleted file mode 100644 index 5415b95a504475517ce0b451729d8edaaf24042a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_cuda_c6404f1 -ops = torch.ops._paged_attention_cuda_c6404f1 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_cuda_c6404f1::{op_name}" diff --git a/build/torch29-cxx11-cu130-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so b/build/torch29-cxx11-cu130-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so deleted file mode 100644 index 841115a9b76b86e9d9e274b6fe42de557182c2c9..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:debaeb20cb67b2c4f707a7f2370400d34175255f07dc548449fe924ead7477d7 -size 86064784 diff --git a/build/torch29-cxx11-cu130-aarch64-linux/metadata.json b/build/torch29-cxx11-cu130-aarch64-linux/metadata.json deleted file mode 100644 index 66651b7d3f95ac9e5ce5fc2a641b6f0f50788f87..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/metadata.json +++ /dev/null @@ -1,19 +0,0 @@ -{ - "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" - ] - } -} diff --git a/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/platforms.py b/build/torch29-cxx11-cu130-aarch64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_custom_ops.py b/build/torch29-cxx11-cu130-x86_64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py b/build/torch29-cxx11-cu130-x86_64-linux/_ops.py deleted file mode 100644 index 5415b95a504475517ce0b451729d8edaaf24042a..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_cuda_c6404f1 -ops = torch.ops._paged_attention_cuda_c6404f1 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_cuda_c6404f1::{op_name}" diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_paged_attention_cuda_c6404f1.abi3.so b/build/torch29-cxx11-cu130-x86_64-linux/_paged_attention_cuda_c6404f1.abi3.so deleted file mode 100644 index fd6bfea111d9609c0fadb2079e517edc916dcf3c..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/_paged_attention_cuda_c6404f1.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a8b109c24c9a2f74e945b6d5bd8634b53b791fbc3ad60bbe7304e2f52097f51e -size 86544216 diff --git a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json b/build/torch29-cxx11-cu130-x86_64-linux/metadata.json deleted file mode 100644 index 66651b7d3f95ac9e5ce5fc2a641b6f0f50788f87..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json +++ /dev/null @@ -1,19 +0,0 @@ -{ - "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" - ] - } -} diff --git a/build/torch29-cxx11-cu130-x86_64-linux/paged_attention/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/platforms.py b/build/torch29-cxx11-cu130-x86_64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/__init__.py b/build/torch29-cxx11-rocm63-x86_64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/_custom_ops.py b/build/torch29-cxx11-rocm63-x86_64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/_ops.py b/build/torch29-cxx11-rocm63-x86_64-linux/_ops.py deleted file mode 100644 index c1268908ac846f37bef7874170ff1f06c3eb33c9..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_rocm_83cf4a3 -ops = torch.ops._paged_attention_rocm_83cf4a3 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_rocm_83cf4a3::{op_name}" diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/_paged_attention_rocm_83cf4a3.abi3.so b/build/torch29-cxx11-rocm63-x86_64-linux/_paged_attention_rocm_83cf4a3.abi3.so deleted file mode 100644 index 639195aac43e3055b6a25ba6e0692c453571e719..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/_paged_attention_rocm_83cf4a3.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ef3e8ae335bb6a87e6f97b0c472cc3aa39c6176240eec858fdcac23b55653118 -size 58641744 diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/metadata.json b/build/torch29-cxx11-rocm63-x86_64-linux/metadata.json deleted file mode 100644 index 3e8d811f1dc42febd33121b2627f809447622baf..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/metadata.json +++ /dev/null @@ -1,17 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "rocm", - "archs": [ - "gfx1030", - "gfx1100", - "gfx1101", - "gfx906", - "gfx908", - "gfx90a", - "gfx942" - ] - } -} diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/paged_attention/__init__.py b/build/torch29-cxx11-rocm63-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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"))) diff --git a/build/torch29-cxx11-rocm63-x86_64-linux/platforms.py b/build/torch29-cxx11-rocm63-x86_64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm63-x86_64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/__init__.py b/build/torch29-cxx11-rocm64-x86_64-linux/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/_custom_ops.py b/build/torch29-cxx11-rocm64-x86_64-linux/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/_ops.py b/build/torch29-cxx11-rocm64-x86_64-linux/_ops.py deleted file mode 100644 index c1268908ac846f37bef7874170ff1f06c3eb33c9..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_rocm_83cf4a3 -ops = torch.ops._paged_attention_rocm_83cf4a3 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_rocm_83cf4a3::{op_name}" diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/_paged_attention_rocm_83cf4a3.abi3.so b/build/torch29-cxx11-rocm64-x86_64-linux/_paged_attention_rocm_83cf4a3.abi3.so deleted file mode 100644 index 9c83dc7ea0dfcf90e995788d714b579c2dc1cdc8..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/_paged_attention_rocm_83cf4a3.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:63144f72f1d8e1d225446f57f89fcdd0e6cce615cf5f6d4c37aa6f33772f292e -size 57972360 diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/metadata.json b/build/torch29-cxx11-rocm64-x86_64-linux/metadata.json deleted file mode 100644 index 3e8d811f1dc42febd33121b2627f809447622baf..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/metadata.json +++ /dev/null @@ -1,17 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [], - "backend": { - "type": "rocm", - "archs": [ - "gfx1030", - "gfx1100", - "gfx1101", - "gfx906", - "gfx908", - "gfx90a", - "gfx942" - ] - } -} diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/paged_attention/__init__.py b/build/torch29-cxx11-rocm64-x86_64-linux/paged_attention/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -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"))) diff --git a/build/torch29-cxx11-rocm64-x86_64-linux/platforms.py b/build/torch29-cxx11-rocm64-x86_64-linux/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-rocm64-x86_64-linux/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -) diff --git a/build/torch29-metal-aarch64-darwin/__init__.py b/build/torch29-metal-aarch64-darwin/__init__.py deleted file mode 100644 index 9de56043369487facc1f163df6bd319c9806e5ca..0000000000000000000000000000000000000000 --- a/build/torch29-metal-aarch64-darwin/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -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", -] diff --git a/build/torch29-metal-aarch64-darwin/_custom_ops.py b/build/torch29-metal-aarch64-darwin/_custom_ops.py deleted file mode 100644 index a0c0b8db085468dee5100c98d14106a9ee917bf2..0000000000000000000000000000000000000000 --- a/build/torch29-metal-aarch64-darwin/_custom_ops.py +++ /dev/null @@ -1,173 +0,0 @@ -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", -] diff --git a/build/torch29-metal-aarch64-darwin/_ops.py b/build/torch29-metal-aarch64-darwin/_ops.py deleted file mode 100644 index b70aaa4f8859a5ed30a12fdf876196d99bea499a..0000000000000000000000000000000000000000 --- a/build/torch29-metal-aarch64-darwin/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _paged_attention_metal_c6404f1 -ops = torch.ops._paged_attention_metal_c6404f1 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_paged_attention_metal_c6404f1::{op_name}" diff --git a/build/torch29-metal-aarch64-darwin/_paged_attention_metal_c6404f1.abi3.so b/build/torch29-metal-aarch64-darwin/_paged_attention_metal_c6404f1.abi3.so deleted file mode 100644 index ed6528f2606d1e91eb6e13f693d22c7c357a962e..0000000000000000000000000000000000000000 --- a/build/torch29-metal-aarch64-darwin/_paged_attention_metal_c6404f1.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:92e72671d13cdb84ccabc5bbc95e4b6e11772708e4e948e6a15974b86e80e26e -size 14893128 diff --git a/build/torch29-metal-aarch64-darwin/metadata.json b/build/torch29-metal-aarch64-darwin/metadata.json deleted file mode 100644 index a5381dd80836f863378b9f33a559815688de9287..0000000000000000000000000000000000000000 --- a/build/torch29-metal-aarch64-darwin/metadata.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "version": 1, - "license": "Apache-2.0", - "python-depends": [] -} \ No newline at end of file diff --git a/build/torch29-metal-aarch64-darwin/paged_attention/__init__.py b/build/torch29-metal-aarch64-darwin/paged_attention/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-metal-aarch64-darwin/paged_attention/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -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"))) diff --git a/build/torch29-metal-aarch64-darwin/platforms.py b/build/torch29-metal-aarch64-darwin/platforms.py deleted file mode 100644 index 6277d5f50ff3ddc265bb39fa1c4d17e0341b7767..0000000000000000000000000000000000000000 --- a/build/torch29-metal-aarch64-darwin/platforms.py +++ /dev/null @@ -1,92 +0,0 @@ -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 -)