Upload 3 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+build/torch27-cxx11-rocm63-x86_64-linux/aiter_cache/_aiter_cache_a120ba4f5566b.so filter=lfs diff=lfs merge=lfs -text

build/torch27-cxx11-rocm63-x86_64-linux/aiter_cache/__init__.py

@@ -0,0 +1,74 @@
+# This is a Python package for the kernel.
+from ._ops import ops
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+from torch import Tensor
+from typing import List, Optional
+from ..jit.core import compile_ops, CK_DIR, AITER_CSRC_DIR
+import torch.nn.functional as F
+
+MD_NAME = "module_cache"
+
+
+def swap_blocks(src: Tensor, dst: Tensor, block_mapping: Tensor):
+    return ops.swap_blocks(src, dst, block_mapping)
+
+def copy_blocks(key_caches: Tensor, value_caches: Tensor, block_mapping: Tensor):
+    return ops.copy_blocks(key_caches, value_caches, block_mapping)
+
+def reshape_and_cache(
+    key: Tensor,
+    value: Tensor,
+    key_cache: Tensor,
+    value_cache: Tensor,
+    slot_mapping: Tensor,
+    kv_cache_dtype: str,
+    k_scale: float,
+    v_scale: float,
+    asm_layout: bool,
+):
+    return ops.reshape_and_cache(
+        key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype, k_scale, v_scale, asm_layout
+    )
+
+def reshape_and_cache_flash(
+    key: Tensor,
+    value: Tensor,
+    key_cache: Tensor,
+    value_cache: Tensor,
+    slot_mapping: Tensor,
+    kv_cache_dtype: str,
+    k_scale: float,
+    v_scale: float,
+):
+    return ops.reshape_and_cache_flash(
+        key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype, k_scale, v_scale
+    )
+
+def reshape_and_cache_with_pertoken_quant(
+    key: Tensor,
+    value: Tensor,
+    key_cache: Tensor,
+    value_cache: Tensor,
+    k_dequant_scales: Tensor,
+    v_dequant_scales: Tensor,
+    slot_mapping: Tensor,
+    asm_layout: bool,
+):
+    return ops.reshape_and_cache_with_pertoken_quant(
+        key, value, key_cache, value_cache, k_dequant_scales, v_dequant_scales, slot_mapping, asm_layout
+    )
+
+def convert_fp8(dst_cache: Tensor, src_cache: Tensor, scale: float, kv_cache_dtype: str):
+    return ops.convert_fp8(dst_cache, src_cache, scale, kv_cache_dtype)
+
+
+__all__ = [
+    "swap_blocks",
+    "copy_blocks",
+    "reshape_and_cache",
+    "reshape_and_cache_flash",
+    "reshape_and_cache_with_pertoken_quant",
+    "convert_fp8",
+]

build/torch27-cxx11-rocm63-x86_64-linux/aiter_cache/_aiter_cache_a120ba4f5566b.so

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

build/torch27-cxx11-rocm63-x86_64-linux/aiter_cache/_ops.py

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