Build uploaded using `kernels`.

build/torch28-cxx11-cu126-x86_64-linux/paged_attention/__init__.py

@@ -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",
-]

build/torch28-cxx11-cu126-x86_64-linux/paged_attention/_custom_ops.py

@@ -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",
-]

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

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

build/torch28-cxx11-cu126-x86_64-linux/paged_attention/_paged_attention_287831d.abi3.so

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

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

@@ -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
-)

build/torch28-cxx11-cu128-x86_64-linux/paged_attention/__init__.py

@@ -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",
-]

build/torch28-cxx11-cu128-x86_64-linux/paged_attention/_custom_ops.py

@@ -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",
-]

build/torch28-cxx11-cu128-x86_64-linux/paged_attention/_ops.py

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

build/torch28-cxx11-cu128-x86_64-linux/paged_attention/_paged_attention_287831d.abi3.so

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

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

@@ -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
-)

build/torch28-cxx11-cu129-x86_64-linux/paged_attention/__init__.py

@@ -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",
-]

build/torch28-cxx11-cu129-x86_64-linux/paged_attention/_custom_ops.py

@@ -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",
-]

build/torch28-cxx11-cu129-x86_64-linux/paged_attention/_ops.py

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

build/torch28-cxx11-cu129-x86_64-linux/paged_attention/_paged_attention_287831d.abi3.so

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

build/torch28-cxx11-cu129-x86_64-linux/paged_attention/platforms.py

@@ -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
-)

build/torch28-cxx11-rocm63-x86_64-linux/paged_attention/__init__.py

@@ -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",
-]

build/torch28-cxx11-rocm63-x86_64-linux/paged_attention/_custom_ops.py

@@ -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",
-]

build/torch28-cxx11-rocm63-x86_64-linux/paged_attention/_ops.py

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

build/torch28-cxx11-rocm63-x86_64-linux/paged_attention/_paged_attention_287831d.abi3.so

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

build/torch28-cxx11-rocm63-x86_64-linux/paged_attention/platforms.py

@@ -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
-)

build/torch28-cxx11-rocm64-x86_64-linux/paged_attention/__init__.py

@@ -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",
-]

build/torch28-cxx11-rocm64-x86_64-linux/paged_attention/_custom_ops.py

@@ -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",
-]

build/torch28-cxx11-rocm64-x86_64-linux/paged_attention/_ops.py

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

build/torch28-cxx11-rocm64-x86_64-linux/paged_attention/_paged_attention_287831d.abi3.so

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

build/torch28-cxx11-rocm64-x86_64-linux/paged_attention/platforms.py

@@ -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
-)

build/torch29-cxx11-cu126-x86_64-linux/paged_attention/__init__.py

@@ -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",
-]

build/torch29-cxx11-cu126-x86_64-linux/paged_attention/_custom_ops.py

@@ -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",
-]

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

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

build/torch29-cxx11-cu126-x86_64-linux/paged_attention/_paged_attention_287831d.abi3.so

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

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

@@ -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
-)

build/torch29-cxx11-cu128-x86_64-linux/paged_attention/__init__.py

@@ -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",
-]

build/torch29-cxx11-cu128-x86_64-linux/paged_attention/_custom_ops.py

@@ -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",
-]

build/torch29-cxx11-cu128-x86_64-linux/paged_attention/_ops.py

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

build/torch29-cxx11-cu128-x86_64-linux/paged_attention/_paged_attention_287831d.abi3.so

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

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

@@ -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
-)

build/torch29-cxx11-cu130-x86_64-linux/paged_attention/__init__.py

@@ -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",
-]

build/torch29-cxx11-cu130-x86_64-linux/paged_attention/_custom_ops.py

@@ -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",
-]

build/torch29-cxx11-cu130-x86_64-linux/paged_attention/_ops.py

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

build/torch29-cxx11-cu130-x86_64-linux/paged_attention/_paged_attention_287831d.abi3.so

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

build/torch29-cxx11-cu130-x86_64-linux/paged_attention/platforms.py

@@ -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
-)

build/torch29-cxx11-rocm63-x86_64-linux/paged_attention/__init__.py

@@ -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",
-]

build/torch29-cxx11-rocm63-x86_64-linux/paged_attention/_custom_ops.py

@@ -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",
-]

build/torch29-cxx11-rocm63-x86_64-linux/paged_attention/_ops.py

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

build/torch29-cxx11-rocm63-x86_64-linux/paged_attention/_paged_attention_287831d.abi3.so

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

build/torch29-cxx11-rocm63-x86_64-linux/paged_attention/platforms.py

@@ -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
-)

build/torch29-cxx11-rocm64-x86_64-linux/paged_attention/__init__.py

@@ -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",
-]

build/torch29-cxx11-rocm64-x86_64-linux/paged_attention/_custom_ops.py

@@ -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",
-]

build/torch29-cxx11-rocm64-x86_64-linux/paged_attention/_ops.py

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

build/torch29-cxx11-rocm64-x86_64-linux/paged_attention/_paged_attention_287831d.abi3.so

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

build/torch29-cxx11-rocm64-x86_64-linux/paged_attention/platforms.py

@@ -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
-)