diff --git a/build/torch210-cxx11-cu126-aarch64-linux/__init__.py b/build/torch210-cxx11-cu126-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9de56043369487facc1f163df6bd319c9806e5ca --- /dev/null +++ b/build/torch210-cxx11-cu126-aarch64-linux/__init__.py @@ -0,0 +1,21 @@ +from ._custom_ops import ( + convert_fp8, + copy_blocks, + paged_attention_v1, + paged_attention_v2, + reshape_and_cache, + reshape_and_cache_flash, + swap_blocks, +) +from ._ops import ops + +__all__ = [ + "convert_fp8", + "copy_blocks", + "ops", + "paged_attention_v1", + "paged_attention_v2", + "reshape_and_cache", + "reshape_and_cache_flash", + "swap_blocks", +] diff --git a/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/_custom_ops.py b/build/torch210-cxx11-cu126-aarch64-linux/_custom_ops.py similarity index 100% rename from build/torch29-cxx11-cu126-aarch64-linux/paged_attention/_custom_ops.py rename to build/torch210-cxx11-cu126-aarch64-linux/_custom_ops.py diff --git a/build/torch210-cxx11-cu126-aarch64-linux/_ops.py b/build/torch210-cxx11-cu126-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..5415b95a504475517ce0b451729d8edaaf24042a --- /dev/null +++ b/build/torch210-cxx11-cu126-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +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/torch210-cxx11-cu126-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so b/build/torch210-cxx11-cu126-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..dbff1ec26b85bfcd5d0dbcdf39b2bd6297131083 --- /dev/null +++ b/build/torch210-cxx11-cu126-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c05dc765dd6a2f6cef9df00b1f167d3dbae9031492b001720d0aba246e48b6b0 +size 140013424 diff --git a/build/torch210-cxx11-cu126-aarch64-linux/metadata.json b/build/torch210-cxx11-cu126-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..f5902b55ab0b2b561c0cf97567c9806c60839c7f --- /dev/null +++ b/build/torch210-cxx11-cu126-aarch64-linux/metadata.json @@ -0,0 +1,18 @@ +{ + "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/torch210-cxx11-cu126-aarch64-linux/paged_attention/__init__.py b/build/torch210-cxx11-cu126-aarch64-linux/paged_attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..03dbc1afe1cf156661a2b1b22003cd5f599a0309 --- /dev/null +++ b/build/torch210-cxx11-cu126-aarch64-linux/paged_attention/__init__.py @@ -0,0 +1,26 @@ +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/paged_attention/platforms.py b/build/torch210-cxx11-cu126-aarch64-linux/platforms.py similarity index 100% rename from build/torch29-cxx11-cu126-aarch64-linux/paged_attention/platforms.py rename to build/torch210-cxx11-cu126-aarch64-linux/platforms.py diff --git a/build/torch210-cxx11-cu128-aarch64-linux/__init__.py b/build/torch210-cxx11-cu128-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9de56043369487facc1f163df6bd319c9806e5ca --- /dev/null +++ b/build/torch210-cxx11-cu128-aarch64-linux/__init__.py @@ -0,0 +1,21 @@ +from ._custom_ops import ( + convert_fp8, + copy_blocks, + paged_attention_v1, + paged_attention_v2, + reshape_and_cache, + reshape_and_cache_flash, + swap_blocks, +) +from ._ops import ops + +__all__ = [ + "convert_fp8", + "copy_blocks", + "ops", + "paged_attention_v1", + "paged_attention_v2", + "reshape_and_cache", + "reshape_and_cache_flash", + "swap_blocks", +] diff --git a/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/_custom_ops.py b/build/torch210-cxx11-cu128-aarch64-linux/_custom_ops.py similarity index 100% rename from build/torch29-cxx11-cu128-aarch64-linux/paged_attention/_custom_ops.py rename to build/torch210-cxx11-cu128-aarch64-linux/_custom_ops.py diff --git a/build/torch210-cxx11-cu128-aarch64-linux/_ops.py b/build/torch210-cxx11-cu128-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..5415b95a504475517ce0b451729d8edaaf24042a --- /dev/null +++ b/build/torch210-cxx11-cu128-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +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/torch210-cxx11-cu128-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so b/build/torch210-cxx11-cu128-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..b8de834319ae7d8138a6782e8301f2228cb15fc3 --- /dev/null +++ b/build/torch210-cxx11-cu128-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:836c3eefb31d02e84812d6432cb57a5092decb3072eed5b1c43af905d2572aae +size 167603936 diff --git a/build/torch210-cxx11-cu128-aarch64-linux/metadata.json b/build/torch210-cxx11-cu128-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..8b796af185fbbd8594fcd846949aa5fadc0ccdda --- /dev/null +++ b/build/torch210-cxx11-cu128-aarch64-linux/metadata.json @@ -0,0 +1,21 @@ +{ + "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/torch210-cxx11-cu128-aarch64-linux/paged_attention/__init__.py b/build/torch210-cxx11-cu128-aarch64-linux/paged_attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..03dbc1afe1cf156661a2b1b22003cd5f599a0309 --- /dev/null +++ b/build/torch210-cxx11-cu128-aarch64-linux/paged_attention/__init__.py @@ -0,0 +1,26 @@ +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/paged_attention/platforms.py b/build/torch210-cxx11-cu128-aarch64-linux/platforms.py similarity index 100% rename from build/torch29-cxx11-cu128-aarch64-linux/paged_attention/platforms.py rename to build/torch210-cxx11-cu128-aarch64-linux/platforms.py diff --git a/build/torch210-cxx11-cu130-aarch64-linux/__init__.py b/build/torch210-cxx11-cu130-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9de56043369487facc1f163df6bd319c9806e5ca --- /dev/null +++ b/build/torch210-cxx11-cu130-aarch64-linux/__init__.py @@ -0,0 +1,21 @@ +from ._custom_ops import ( + convert_fp8, + copy_blocks, + paged_attention_v1, + paged_attention_v2, + reshape_and_cache, + reshape_and_cache_flash, + swap_blocks, +) +from ._ops import ops + +__all__ = [ + "convert_fp8", + "copy_blocks", + "ops", + "paged_attention_v1", + "paged_attention_v2", + "reshape_and_cache", + "reshape_and_cache_flash", + "swap_blocks", +] diff --git a/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/_custom_ops.py b/build/torch210-cxx11-cu130-aarch64-linux/_custom_ops.py similarity index 100% rename from build/torch29-cxx11-cu130-aarch64-linux/paged_attention/_custom_ops.py rename to build/torch210-cxx11-cu130-aarch64-linux/_custom_ops.py diff --git a/build/torch210-cxx11-cu130-aarch64-linux/_ops.py b/build/torch210-cxx11-cu130-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..5415b95a504475517ce0b451729d8edaaf24042a --- /dev/null +++ b/build/torch210-cxx11-cu130-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +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/torch210-cxx11-cu130-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so b/build/torch210-cxx11-cu130-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..7a80831ae611541ab238c0d8b0f3749665c81d51 --- /dev/null +++ b/build/torch210-cxx11-cu130-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3111815e8cea4da876db81e76a2f14076ff65a35dedb31049bd1be19d966da2f +size 86068816 diff --git a/build/torch210-cxx11-cu130-aarch64-linux/metadata.json b/build/torch210-cxx11-cu130-aarch64-linux/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..66651b7d3f95ac9e5ce5fc2a641b6f0f50788f87 --- /dev/null +++ b/build/torch210-cxx11-cu130-aarch64-linux/metadata.json @@ -0,0 +1,19 @@ +{ + "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/torch210-cxx11-cu130-aarch64-linux/paged_attention/__init__.py b/build/torch210-cxx11-cu130-aarch64-linux/paged_attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..03dbc1afe1cf156661a2b1b22003cd5f599a0309 --- /dev/null +++ b/build/torch210-cxx11-cu130-aarch64-linux/paged_attention/__init__.py @@ -0,0 +1,26 @@ +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/paged_attention/platforms.py b/build/torch210-cxx11-cu130-aarch64-linux/platforms.py similarity index 100% rename from build/torch29-cxx11-cu130-aarch64-linux/paged_attention/platforms.py rename to build/torch210-cxx11-cu130-aarch64-linux/platforms.py diff --git a/build/torch29-cxx11-cu126-aarch64-linux/__init__.py b/build/torch29-cxx11-cu126-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9de56043369487facc1f163df6bd319c9806e5ca --- /dev/null +++ b/build/torch29-cxx11-cu126-aarch64-linux/__init__.py @@ -0,0 +1,21 @@ +from ._custom_ops import ( + convert_fp8, + copy_blocks, + paged_attention_v1, + paged_attention_v2, + reshape_and_cache, + reshape_and_cache_flash, + swap_blocks, +) +from ._ops import ops + +__all__ = [ + "convert_fp8", + "copy_blocks", + "ops", + "paged_attention_v1", + "paged_attention_v2", + "reshape_and_cache", + "reshape_and_cache_flash", + "swap_blocks", +] diff --git a/build/torch29-cxx11-cu126-aarch64-linux/_custom_ops.py b/build/torch29-cxx11-cu126-aarch64-linux/_custom_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..a0c0b8db085468dee5100c98d14106a9ee917bf2 --- /dev/null +++ b/build/torch29-cxx11-cu126-aarch64-linux/_custom_ops.py @@ -0,0 +1,173 @@ +from typing import List, Optional + +import torch + +from ._ops import ops + + +# page attention ops +def paged_attention_v1( + out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_seq_len: int, + alibi_slopes: Optional[torch.Tensor], + kv_cache_dtype: str, + k_scale: float, + v_scale: float, + tp_rank: int = 0, + blocksparse_local_blocks: int = 0, + blocksparse_vert_stride: int = 0, + blocksparse_block_size: int = 64, + blocksparse_head_sliding_step: int = 0, +) -> None: + ops.paged_attention_v1( + out, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + tp_rank, + blocksparse_local_blocks, + blocksparse_vert_stride, + blocksparse_block_size, + blocksparse_head_sliding_step, + ) + + +def paged_attention_v2( + out: torch.Tensor, + exp_sum: torch.Tensor, + max_logits: torch.Tensor, + tmp_out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_seq_len: int, + alibi_slopes: Optional[torch.Tensor], + kv_cache_dtype: str, + k_scale: float, + v_scale: float, + tp_rank: int = 0, + blocksparse_local_blocks: int = 0, + blocksparse_vert_stride: int = 0, + blocksparse_block_size: int = 64, + blocksparse_head_sliding_step: int = 0, +) -> None: + ops.paged_attention_v2( + out, + exp_sum, + max_logits, + tmp_out, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + tp_rank, + blocksparse_local_blocks, + blocksparse_vert_stride, + blocksparse_block_size, + blocksparse_head_sliding_step, + ) + + +def reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: float, + v_scale: float, +) -> None: + ops.reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + + +def reshape_and_cache_flash( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, +) -> None: + ops.reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + + +def copy_blocks( + key_caches: List[torch.Tensor], + value_caches: List[torch.Tensor], + block_mapping: torch.Tensor, +) -> None: + ops.copy_blocks(key_caches, value_caches, block_mapping) + + +def swap_blocks( + src: torch.Tensor, dst: torch.Tensor, block_mapping: torch.Tensor +) -> None: + ops.swap_blocks(src, dst, block_mapping) + + +def convert_fp8( + output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8" +) -> None: + ops.convert_fp8(output, input, scale, kv_dtype) + + +__all__ = [ + "convert_fp8", + "paged_attention_v1", + "paged_attention_v2", + "reshape_and_cache", + "copy_blocks", +] diff --git a/build/torch29-cxx11-cu126-aarch64-linux/_ops.py b/build/torch29-cxx11-cu126-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..5415b95a504475517ce0b451729d8edaaf24042a --- /dev/null +++ b/build/torch29-cxx11-cu126-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..1760b475fbbb0d8a67982aaccbd8d83fc598157c --- /dev/null +++ b/build/torch29-cxx11-cu126-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so @@ -0,0 +1,3 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..f5902b55ab0b2b561c0cf97567c9806c60839c7f --- /dev/null +++ b/build/torch29-cxx11-cu126-aarch64-linux/metadata.json @@ -0,0 +1,18 @@ +{ + "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 index 9de56043369487facc1f163df6bd319c9806e5ca..03dbc1afe1cf156661a2b1b22003cd5f599a0309 100644 --- a/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__init__.py +++ b/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__init__.py @@ -1,21 +1,26 @@ -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 +import ctypes +import sys -__all__ = [ - "convert_fp8", - "copy_blocks", - "ops", - "paged_attention_v1", - "paged_attention_v2", - "reshape_and_cache", - "reshape_and_cache_flash", - "swap_blocks", -] +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/paged_attention/__pycache__/__init__.cpython-313.pyc b/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index 4c35d9c01a28b3d52578f20c0116c30fd3aee5e7..0000000000000000000000000000000000000000 Binary files a/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc b/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc deleted file mode 100644 index a7ee2b17c2afd046a6634541c2a78148adc920b4..0000000000000000000000000000000000000000 Binary files a/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc b/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index b05df5ac67a38cdaee7a1c65ae83b0df7476efe8..0000000000000000000000000000000000000000 Binary files a/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/_ops.py b/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/_ops.py deleted file mode 100644 index a5e9f56ff02bdcf047fd97ef2a46b10d6fb5e5eb..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-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/torch29-cxx11-cu126-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so b/build/torch29-cxx11-cu126-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so deleted file mode 100755 index e147e4816316407f5430e44670a08020a4d2d1b6..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-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:75a193f66139a77108969642a435eb345b97c24e4a31e53856fd7c65bff6e810 -size 110584256 diff --git a/build/torch29-cxx11-cu126-aarch64-linux/platforms.py b/build/torch29-cxx11-cu126-aarch64-linux/platforms.py new file mode 100644 index 0000000000000000000000000000000000000000..6277d5f50ff3ddc265bb39fa1c4d17e0341b7767 --- /dev/null +++ b/build/torch29-cxx11-cu126-aarch64-linux/platforms.py @@ -0,0 +1,92 @@ +import os +import random +from abc import ABC, abstractmethod +from functools import lru_cache, wraps +from typing import Callable, ParamSpec, TypeVar + +import numpy as np +import torch + +IS_ROCM = torch.version.hip is not None +IS_MPS = torch.backends.mps.is_available() + + +class Platform(ABC): + @classmethod + def seed_everything(cls, seed: int) -> None: + """ + Set the seed of each random module. + `torch.manual_seed` will set seed on all devices. + + Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20 + """ + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + @abstractmethod + def get_device_name(self, device_id: int = 0) -> str: ... + + @abstractmethod + def is_cuda(self) -> bool: ... + + @abstractmethod + def is_rocm(self) -> bool: ... + + @abstractmethod + def is_mps(self) -> bool: ... + + +class CudaPlatform(Platform): + @classmethod + @lru_cache(maxsize=8) + def get_device_name(cls, device_id: int = 0) -> str: + return torch.cuda.get_device_name(0) + + def is_cuda(self) -> bool: + return True + + def is_rocm(self) -> bool: + return False + + def is_mps(self) -> bool: + return False + + +class RocmPlatform(Platform): + @classmethod + @lru_cache(maxsize=8) + def get_device_name(cls, device_id: int = 0) -> str: + return torch.cuda.get_device_name(device_id) + + def is_cuda(self) -> bool: + return False + + def is_rocm(self) -> bool: + return True + + def is_mps(self) -> bool: + return False + + +class MpsPlatform(Platform): + @classmethod + @lru_cache(maxsize=8) + def get_device_name(cls, device_id: int = 0) -> str: + return torch.cuda.get_device_name(device_id) + + def is_cuda(self) -> bool: + return False + + def is_rocm(self) -> bool: + return False + + def is_mps(self) -> bool: + return True + +current_platform = ( + RocmPlatform() if IS_ROCM else + MpsPlatform() if IS_MPS else + CudaPlatform() if torch.cuda.is_available() else + None +) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9de56043369487facc1f163df6bd319c9806e5ca --- /dev/null +++ b/build/torch29-cxx11-cu128-aarch64-linux/__init__.py @@ -0,0 +1,21 @@ +from ._custom_ops import ( + convert_fp8, + copy_blocks, + paged_attention_v1, + paged_attention_v2, + reshape_and_cache, + reshape_and_cache_flash, + swap_blocks, +) +from ._ops import ops + +__all__ = [ + "convert_fp8", + "copy_blocks", + "ops", + "paged_attention_v1", + "paged_attention_v2", + "reshape_and_cache", + "reshape_and_cache_flash", + "swap_blocks", +] diff --git a/build/torch29-cxx11-cu128-aarch64-linux/_custom_ops.py b/build/torch29-cxx11-cu128-aarch64-linux/_custom_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..a0c0b8db085468dee5100c98d14106a9ee917bf2 --- /dev/null +++ b/build/torch29-cxx11-cu128-aarch64-linux/_custom_ops.py @@ -0,0 +1,173 @@ +from typing import List, Optional + +import torch + +from ._ops import ops + + +# page attention ops +def paged_attention_v1( + out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_seq_len: int, + alibi_slopes: Optional[torch.Tensor], + kv_cache_dtype: str, + k_scale: float, + v_scale: float, + tp_rank: int = 0, + blocksparse_local_blocks: int = 0, + blocksparse_vert_stride: int = 0, + blocksparse_block_size: int = 64, + blocksparse_head_sliding_step: int = 0, +) -> None: + ops.paged_attention_v1( + out, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + tp_rank, + blocksparse_local_blocks, + blocksparse_vert_stride, + blocksparse_block_size, + blocksparse_head_sliding_step, + ) + + +def paged_attention_v2( + out: torch.Tensor, + exp_sum: torch.Tensor, + max_logits: torch.Tensor, + tmp_out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_seq_len: int, + alibi_slopes: Optional[torch.Tensor], + kv_cache_dtype: str, + k_scale: float, + v_scale: float, + tp_rank: int = 0, + blocksparse_local_blocks: int = 0, + blocksparse_vert_stride: int = 0, + blocksparse_block_size: int = 64, + blocksparse_head_sliding_step: int = 0, +) -> None: + ops.paged_attention_v2( + out, + exp_sum, + max_logits, + tmp_out, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + tp_rank, + blocksparse_local_blocks, + blocksparse_vert_stride, + blocksparse_block_size, + blocksparse_head_sliding_step, + ) + + +def reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: float, + v_scale: float, +) -> None: + ops.reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + + +def reshape_and_cache_flash( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, +) -> None: + ops.reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + + +def copy_blocks( + key_caches: List[torch.Tensor], + value_caches: List[torch.Tensor], + block_mapping: torch.Tensor, +) -> None: + ops.copy_blocks(key_caches, value_caches, block_mapping) + + +def swap_blocks( + src: torch.Tensor, dst: torch.Tensor, block_mapping: torch.Tensor +) -> None: + ops.swap_blocks(src, dst, block_mapping) + + +def convert_fp8( + output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8" +) -> None: + ops.convert_fp8(output, input, scale, kv_dtype) + + +__all__ = [ + "convert_fp8", + "paged_attention_v1", + "paged_attention_v2", + "reshape_and_cache", + "copy_blocks", +] diff --git a/build/torch29-cxx11-cu128-aarch64-linux/_ops.py b/build/torch29-cxx11-cu128-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..5415b95a504475517ce0b451729d8edaaf24042a --- /dev/null +++ b/build/torch29-cxx11-cu128-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..e4641642e57ed6726f75d4e867b337f0f01ad4f8 --- /dev/null +++ b/build/torch29-cxx11-cu128-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so @@ -0,0 +1,3 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..8b796af185fbbd8594fcd846949aa5fadc0ccdda --- /dev/null +++ b/build/torch29-cxx11-cu128-aarch64-linux/metadata.json @@ -0,0 +1,21 @@ +{ + "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 index 9de56043369487facc1f163df6bd319c9806e5ca..03dbc1afe1cf156661a2b1b22003cd5f599a0309 100644 --- a/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__init__.py +++ b/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__init__.py @@ -1,21 +1,26 @@ -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 +import ctypes +import sys -__all__ = [ - "convert_fp8", - "copy_blocks", - "ops", - "paged_attention_v1", - "paged_attention_v2", - "reshape_and_cache", - "reshape_and_cache_flash", - "swap_blocks", -] +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/paged_attention/__pycache__/__init__.cpython-313.pyc b/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index d9b411f561235eeebf4831cc5b7fd3af6d6b2b97..0000000000000000000000000000000000000000 Binary files a/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc b/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc deleted file mode 100644 index 8eb24f6e486563cd362705980c39de6d933f9e49..0000000000000000000000000000000000000000 Binary files a/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc b/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index 89df1e0361267bd24e3d4341e77b75daf8a187a9..0000000000000000000000000000000000000000 Binary files a/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/_ops.py b/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/_ops.py deleted file mode 100644 index a5e9f56ff02bdcf047fd97ef2a46b10d6fb5e5eb..0000000000000000000000000000000000000000 --- a/build/torch29-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/torch29-cxx11-cu128-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so b/build/torch29-cxx11-cu128-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so deleted file mode 100755 index d514c5d024e62151a13638101a3c96d23b0ee198..0000000000000000000000000000000000000000 --- a/build/torch29-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:f0955cb5ea9b9714579c81efc8810bba7ce46aad33701c244efcb01d021a944f -size 138174216 diff --git a/build/torch29-cxx11-cu128-aarch64-linux/platforms.py b/build/torch29-cxx11-cu128-aarch64-linux/platforms.py new file mode 100644 index 0000000000000000000000000000000000000000..6277d5f50ff3ddc265bb39fa1c4d17e0341b7767 --- /dev/null +++ b/build/torch29-cxx11-cu128-aarch64-linux/platforms.py @@ -0,0 +1,92 @@ +import os +import random +from abc import ABC, abstractmethod +from functools import lru_cache, wraps +from typing import Callable, ParamSpec, TypeVar + +import numpy as np +import torch + +IS_ROCM = torch.version.hip is not None +IS_MPS = torch.backends.mps.is_available() + + +class Platform(ABC): + @classmethod + def seed_everything(cls, seed: int) -> None: + """ + Set the seed of each random module. + `torch.manual_seed` will set seed on all devices. + + Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20 + """ + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + @abstractmethod + def get_device_name(self, device_id: int = 0) -> str: ... + + @abstractmethod + def is_cuda(self) -> bool: ... + + @abstractmethod + def is_rocm(self) -> bool: ... + + @abstractmethod + def is_mps(self) -> bool: ... + + +class CudaPlatform(Platform): + @classmethod + @lru_cache(maxsize=8) + def get_device_name(cls, device_id: int = 0) -> str: + return torch.cuda.get_device_name(0) + + def is_cuda(self) -> bool: + return True + + def is_rocm(self) -> bool: + return False + + def is_mps(self) -> bool: + return False + + +class RocmPlatform(Platform): + @classmethod + @lru_cache(maxsize=8) + def get_device_name(cls, device_id: int = 0) -> str: + return torch.cuda.get_device_name(device_id) + + def is_cuda(self) -> bool: + return False + + def is_rocm(self) -> bool: + return True + + def is_mps(self) -> bool: + return False + + +class MpsPlatform(Platform): + @classmethod + @lru_cache(maxsize=8) + def get_device_name(cls, device_id: int = 0) -> str: + return torch.cuda.get_device_name(device_id) + + def is_cuda(self) -> bool: + return False + + def is_rocm(self) -> bool: + return False + + def is_mps(self) -> bool: + return True + +current_platform = ( + RocmPlatform() if IS_ROCM else + MpsPlatform() if IS_MPS else + CudaPlatform() if torch.cuda.is_available() else + None +) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9de56043369487facc1f163df6bd319c9806e5ca --- /dev/null +++ b/build/torch29-cxx11-cu130-aarch64-linux/__init__.py @@ -0,0 +1,21 @@ +from ._custom_ops import ( + convert_fp8, + copy_blocks, + paged_attention_v1, + paged_attention_v2, + reshape_and_cache, + reshape_and_cache_flash, + swap_blocks, +) +from ._ops import ops + +__all__ = [ + "convert_fp8", + "copy_blocks", + "ops", + "paged_attention_v1", + "paged_attention_v2", + "reshape_and_cache", + "reshape_and_cache_flash", + "swap_blocks", +] diff --git a/build/torch29-cxx11-cu130-aarch64-linux/_custom_ops.py b/build/torch29-cxx11-cu130-aarch64-linux/_custom_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..a0c0b8db085468dee5100c98d14106a9ee917bf2 --- /dev/null +++ b/build/torch29-cxx11-cu130-aarch64-linux/_custom_ops.py @@ -0,0 +1,173 @@ +from typing import List, Optional + +import torch + +from ._ops import ops + + +# page attention ops +def paged_attention_v1( + out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_seq_len: int, + alibi_slopes: Optional[torch.Tensor], + kv_cache_dtype: str, + k_scale: float, + v_scale: float, + tp_rank: int = 0, + blocksparse_local_blocks: int = 0, + blocksparse_vert_stride: int = 0, + blocksparse_block_size: int = 64, + blocksparse_head_sliding_step: int = 0, +) -> None: + ops.paged_attention_v1( + out, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + tp_rank, + blocksparse_local_blocks, + blocksparse_vert_stride, + blocksparse_block_size, + blocksparse_head_sliding_step, + ) + + +def paged_attention_v2( + out: torch.Tensor, + exp_sum: torch.Tensor, + max_logits: torch.Tensor, + tmp_out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_seq_len: int, + alibi_slopes: Optional[torch.Tensor], + kv_cache_dtype: str, + k_scale: float, + v_scale: float, + tp_rank: int = 0, + blocksparse_local_blocks: int = 0, + blocksparse_vert_stride: int = 0, + blocksparse_block_size: int = 64, + blocksparse_head_sliding_step: int = 0, +) -> None: + ops.paged_attention_v2( + out, + exp_sum, + max_logits, + tmp_out, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + tp_rank, + blocksparse_local_blocks, + blocksparse_vert_stride, + blocksparse_block_size, + blocksparse_head_sliding_step, + ) + + +def reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: float, + v_scale: float, +) -> None: + ops.reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + + +def reshape_and_cache_flash( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + v_scale: torch.Tensor, +) -> None: + ops.reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + + +def copy_blocks( + key_caches: List[torch.Tensor], + value_caches: List[torch.Tensor], + block_mapping: torch.Tensor, +) -> None: + ops.copy_blocks(key_caches, value_caches, block_mapping) + + +def swap_blocks( + src: torch.Tensor, dst: torch.Tensor, block_mapping: torch.Tensor +) -> None: + ops.swap_blocks(src, dst, block_mapping) + + +def convert_fp8( + output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8" +) -> None: + ops.convert_fp8(output, input, scale, kv_dtype) + + +__all__ = [ + "convert_fp8", + "paged_attention_v1", + "paged_attention_v2", + "reshape_and_cache", + "copy_blocks", +] diff --git a/build/torch29-cxx11-cu130-aarch64-linux/_ops.py b/build/torch29-cxx11-cu130-aarch64-linux/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..5415b95a504475517ce0b451729d8edaaf24042a --- /dev/null +++ b/build/torch29-cxx11-cu130-aarch64-linux/_ops.py @@ -0,0 +1,9 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..841115a9b76b86e9d9e274b6fe42de557182c2c9 --- /dev/null +++ b/build/torch29-cxx11-cu130-aarch64-linux/_paged_attention_cuda_c6404f1.abi3.so @@ -0,0 +1,3 @@ +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 new file mode 100644 index 0000000000000000000000000000000000000000..66651b7d3f95ac9e5ce5fc2a641b6f0f50788f87 --- /dev/null +++ b/build/torch29-cxx11-cu130-aarch64-linux/metadata.json @@ -0,0 +1,19 @@ +{ + "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 index 9de56043369487facc1f163df6bd319c9806e5ca..03dbc1afe1cf156661a2b1b22003cd5f599a0309 100644 --- a/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__init__.py +++ b/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__init__.py @@ -1,21 +1,26 @@ -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 +import ctypes +import sys -__all__ = [ - "convert_fp8", - "copy_blocks", - "ops", - "paged_attention_v1", - "paged_attention_v2", - "reshape_and_cache", - "reshape_and_cache_flash", - "swap_blocks", -] +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/paged_attention/__pycache__/__init__.cpython-313.pyc b/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index f8fb851276be004995178c84b861a35401203ccc..0000000000000000000000000000000000000000 Binary files a/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc b/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc deleted file mode 100644 index 3f7ec2073f8594850264af2ff274d089067cc2b0..0000000000000000000000000000000000000000 Binary files a/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__pycache__/_custom_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc b/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc deleted file mode 100644 index 389abfdad4679d365116f9341819ed8d89fd1e8a..0000000000000000000000000000000000000000 Binary files a/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/__pycache__/_ops.cpython-313.pyc and /dev/null differ diff --git a/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/_ops.py b/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/_ops.py deleted file mode 100644 index a5e9f56ff02bdcf047fd97ef2a46b10d6fb5e5eb..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-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/torch29-cxx11-cu130-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so b/build/torch29-cxx11-cu130-aarch64-linux/paged_attention/_paged_attention_b4c51e9.abi3.so deleted file mode 100755 index 9fe5180086bb402126f2bd617343f8814eecfe2c..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-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:6515c6d2e299021620445210f11a918d6b8cd93c8f8912b6904b40a86d646561 -size 76628152 diff --git a/build/torch29-cxx11-cu130-aarch64-linux/platforms.py b/build/torch29-cxx11-cu130-aarch64-linux/platforms.py new file mode 100644 index 0000000000000000000000000000000000000000..6277d5f50ff3ddc265bb39fa1c4d17e0341b7767 --- /dev/null +++ b/build/torch29-cxx11-cu130-aarch64-linux/platforms.py @@ -0,0 +1,92 @@ +import os +import random +from abc import ABC, abstractmethod +from functools import lru_cache, wraps +from typing import Callable, ParamSpec, TypeVar + +import numpy as np +import torch + +IS_ROCM = torch.version.hip is not None +IS_MPS = torch.backends.mps.is_available() + + +class Platform(ABC): + @classmethod + def seed_everything(cls, seed: int) -> None: + """ + Set the seed of each random module. + `torch.manual_seed` will set seed on all devices. + + Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20 + """ + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + @abstractmethod + def get_device_name(self, device_id: int = 0) -> str: ... + + @abstractmethod + def is_cuda(self) -> bool: ... + + @abstractmethod + def is_rocm(self) -> bool: ... + + @abstractmethod + def is_mps(self) -> bool: ... + + +class CudaPlatform(Platform): + @classmethod + @lru_cache(maxsize=8) + def get_device_name(cls, device_id: int = 0) -> str: + return torch.cuda.get_device_name(0) + + def is_cuda(self) -> bool: + return True + + def is_rocm(self) -> bool: + return False + + def is_mps(self) -> bool: + return False + + +class RocmPlatform(Platform): + @classmethod + @lru_cache(maxsize=8) + def get_device_name(cls, device_id: int = 0) -> str: + return torch.cuda.get_device_name(device_id) + + def is_cuda(self) -> bool: + return False + + def is_rocm(self) -> bool: + return True + + def is_mps(self) -> bool: + return False + + +class MpsPlatform(Platform): + @classmethod + @lru_cache(maxsize=8) + def get_device_name(cls, device_id: int = 0) -> str: + return torch.cuda.get_device_name(device_id) + + def is_cuda(self) -> bool: + return False + + def is_rocm(self) -> bool: + return False + + def is_mps(self) -> bool: + return True + +current_platform = ( + RocmPlatform() if IS_ROCM else + MpsPlatform() if IS_MPS else + CudaPlatform() if torch.cuda.is_available() else + None +)