tileops 0.0.1.dev0__tar.gz

tileops-0.0.1.dev0/LICENSE

@@ -0,0 +1,21 @@
+    MIT License
+
+    Copyright (c) Tile-AI.
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE

tileops-0.0.1.dev0/PKG-INFO

@@ -0,0 +1,82 @@
+Metadata-Version: 2.4
+Name: tileops
+Version: 0.0.1.dev0
+Summary: TileOPs kernels for efficient inference
+Author: Tile-AI
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/tile-ai/tileops
+Project-URL: Bug Tracker, https://github.com/tile-ai/tileops/issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: tilelang
+Dynamic: license-file
+
+# TileOPs (TOP)
+
+**TileOPs (TOP)** is a high-performance machine learning operator collections built to run on the [TileLang](https://github.com/tile-ai/tilelang) backend. It offers efficient, modular, and composable implementations optimized for AI workloads.
+
+---
+
+
+## 📦 Installation
+
+### Requirements
+
+- Python 3.8+
+- PyTorch >= 2.1
+- [TileLang](https://github.com/tilelang/tilelang)
+- Triton (optional, for selected fast kernels)
+
+### Install (editable mode for development)
+
+```bash
+git clone https://github.com/tile-ai/TileOPs
+cd TileOPs
+git submodule update --init --recursive
+pip install -e .
+```
+
+## 🚀 Quick Usage
+
+```python
+import torch
+import top
+from top import MLAKernel
+
+device = "cuda"
+dtype = torch.float16
+
+batch = 128
+heads = 64
+kv_heads = 1
+kv_ctx = 8192
+dim = 512
+pe_dim = 64
+
+# Query input: [batch, heads, dim]
+q = torch.randn(batch, heads, dim, device=device, dtype=dtype)
+
+# Query positional encoding: [batch, heads, pe_dim]
+q_pe = torch.randn(batch, heads, pe_dim, device=device, dtype=dtype)
+
+# KV cache input: [batch, kv_ctx, kv_heads, dim]
+kv = torch.randn(batch, kv_ctx, kv_heads, dim, device=device, dtype=dtype)
+
+# KV positional encoding: [batch, kv_ctx, kv_heads, pe_dim]
+k_pe = torch.randn(batch, kv_ctx, kv_heads, pe_dim, device=device, dtype=dtype)
+
+# Use MLA kernel
+block_N = 64
+block_H = 64
+num_split = 1
+
+mla = MLAKernel(batch, heads, kv_heads, kv_ctx, dim, pe_dim, block_N, block_H, num_split)
+
+out = mla(q, q_pe, kv, k_pe)
+```
+
+## Acknowledgments

tileops-0.0.1.dev0/README.md

@@ -0,0 +1,65 @@
+# TileOPs (TOP)
+
+**TileOPs (TOP)** is a high-performance machine learning operator collections built to run on the [TileLang](https://github.com/tile-ai/tilelang) backend. It offers efficient, modular, and composable implementations optimized for AI workloads.
+
+---
+
+
+## 📦 Installation
+
+### Requirements
+
+- Python 3.8+
+- PyTorch >= 2.1
+- [TileLang](https://github.com/tilelang/tilelang)
+- Triton (optional, for selected fast kernels)
+
+### Install (editable mode for development)
+
+```bash
+git clone https://github.com/tile-ai/TileOPs
+cd TileOPs
+git submodule update --init --recursive
+pip install -e .
+```
+
+## 🚀 Quick Usage
+
+```python
+import torch
+import top
+from top import MLAKernel
+
+device = "cuda"
+dtype = torch.float16
+
+batch = 128
+heads = 64
+kv_heads = 1
+kv_ctx = 8192
+dim = 512
+pe_dim = 64
+
+# Query input: [batch, heads, dim]
+q = torch.randn(batch, heads, dim, device=device, dtype=dtype)
+
+# Query positional encoding: [batch, heads, pe_dim]
+q_pe = torch.randn(batch, heads, pe_dim, device=device, dtype=dtype)
+
+# KV cache input: [batch, kv_ctx, kv_heads, dim]
+kv = torch.randn(batch, kv_ctx, kv_heads, dim, device=device, dtype=dtype)
+
+# KV positional encoding: [batch, kv_ctx, kv_heads, pe_dim]
+k_pe = torch.randn(batch, kv_ctx, kv_heads, pe_dim, device=device, dtype=dtype)
+
+# Use MLA kernel
+block_N = 64
+block_H = 64
+num_split = 1
+
+mla = MLAKernel(batch, heads, kv_heads, kv_ctx, dim, pe_dim, block_N, block_H, num_split)
+
+out = mla(q, q_pe, kv, k_pe)
+```
+
+## Acknowledgments

tileops-0.0.1.dev0/pyproject.toml

@@ -0,0 +1,70 @@
+[project]
+name = "tileops"
+version = "0.0.1.dev0"
+description = "TileOPs kernels for efficient inference"
+authors = [{ name="Tile-AI" }]
+readme = "README.md"
+requires-python = ">=3.8"
+license = "MIT"
+dependencies = ["torch", "tilelang"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Programming Language :: Python :: 3"
+]
+
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.yapf]
+based_on_style = "yapf"
+column_limit = 100
+indent_width = 4
+
+[tool.codespell]
+ignore-words-list = "nd, te, ist, LOD, offen, NotIn, HSA"
+skip = [
+    "build",
+    "3rdparty",
+    "dist",
+    ".venv"
+]
+
+[tool.setuptools]
+packages = ["top"]
+
+[project.urls]
+"Homepage" = "https://github.com/tile-ai/tileops"
+"Bug Tracker" = "https://github.com/tile-ai/tileops/issues"
+
+[tool.ruff.lint]
+select = [
+    # pycodestyle
+    "E",
+    # Pyflakes
+    "F",
+    # pyupgrade
+    # "UP",
+    # flake8-bugbear
+    "B",
+    # flake8-simplify
+    "SIM",
+    # isort
+    # "I",
+]
+ignore = [
+    # Module level import not at top of file
+    "E402",
+    # star imports
+    "F405", "F403",
+    # ambiguous name
+    "E741",
+    # line too long
+    "E501",
+    # key in dict.keys()
+    "SIM118",
+    # memory leaks
+    "B019",
+    # No such file or directory
+    "E902",
+]

tileops-0.0.1.dev0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

tileops-0.0.1.dev0/setup.py

@@ -0,0 +1,33 @@
+# Copyright (c) Tile-AI.
+# Licensed under the MIT License.
+
+import subprocess
+from setuptools import setup, find_packages
+from setuptools.command.install import install
+
+
+class CustomInstall(install):
+
+    def run(self):
+        # Step 1: Install system dependencies (optional, warning: requires sudo)
+        # print("Installing system dependencies for TileLang...")
+        # subprocess.run(["sudo", "apt-get", "update"])
+        # subprocess.run([
+        #     "sudo", "apt-get", "install", "-y", "python3-setuptools", "gcc",
+        #     "libtinfo-dev", "zlib1g-dev", "build-essential", "cmake",
+        #     "libedit-dev", "libxml2-dev"
+        # ])
+
+        # Step 2: Install TileLang via pip from local path
+        # print("Installing TileLang from submodule...")
+        # subprocess.run(["pip", "install", "-e", "3rdparty/tilelang"])
+
+        # Continue normal installation
+        install.run(self)
+
+
+setup(name="top",
+      version="0.1.0",
+      packages=find_packages(),
+      install_requires=["torch", "tilelang"],
+      cmdclass={'install': CustomInstall})

tileops-0.0.1.dev0/tests/test_bitnet.py

@@ -0,0 +1,34 @@
+import argparse
+import torch
+from top import Bitnet_158_int8xint2_kernel
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--M', type=int, default=1, help='')
+    parser.add_argument('--N', type=int, default=256, help='')
+    parser.add_argument('--K', type=int, default=256, help='')
+    parser.add_argument('--in_dtype', type=str, default="int8", help='')
+    parser.add_argument('--out_dtype', type=str, default="int32", help='')
+    parser.add_argument('--accum_dtype', type=str, default="int32", help='')
+    parser.add_argument('--fast_decoding', type=bool, default=True, help='')
+
+    args = parser.parse_args()
+    M, N, K, in_dtype, out_dtype, accum_dtype, fast_decoding = args.M, args.N, args.K, args.in_dtype, args.out_dtype, args.accum_dtype, args.fast_decoding
+
+    A = torch.randint(0, 4, (M, K), device="cuda", dtype=getattr(torch, in_dtype))
+    B = torch.randint(0, 2, (N, K), device="cuda", dtype=getattr(torch, in_dtype))
+    C = torch.zeros(M, N, device="cuda", dtype=getattr(torch, accum_dtype))
+
+    bitnet_158_int8xint2 = Bitnet_158_int8xint2_kernel(
+        M, N, K, in_dtype, out_dtype, accum_dtype, fast_decoding)
+    C1 = bitnet_158_int8xint2.prefill(A, B, C)
+    C2 = bitnet_158_int8xint2.decode(A, B, C)
+    print(C1)
+    print(C2)
+    bitnet_158_int8xint2.check(A, B, C, "prefill")
+    bitnet_158_int8xint2.check(A, B, C, "decode")
+
+
+if __name__ == "__main__":
+    main()

tileops-0.0.1.dev0/tests/test_blocksparse_attention.py

@@ -0,0 +1,55 @@
+# Copyright (c) Tile-AI.
+# Licensed under the MIT License.
+
+import argparse
+import torch
+from top import BlockSparseAttentionKernel
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=8, help='batch size')
+    parser.add_argument('--h', type=int, default=32, help='Number of heads')
+    parser.add_argument('--n_ctx', type=int, default=1024, help='Context size')
+    parser.add_argument('--d_head_qk', type=int, default=128, help='Head dimension for Q/K')
+    parser.add_argument('--d_head_v', type=int, default=64, help='Head dimension for V')
+    parser.add_argument('--causal', action='store_true', help='Causal flag')
+    parser.add_argument('--groups', type=int, default=16, help='groups')
+    args = parser.parse_args()
+    BATCH, H, N_CTX, D_HEAD_QK, D_HEAD_V, groups, causal = args.batch, args.h, args.n_ctx, args.d_head_qk, args.d_head_v, args.groups, args.causal
+
+    BLOCK_M = 128
+    BLOCK_N = 64
+    Q = (
+        torch.empty(BATCH, N_CTX, H, D_HEAD_QK, dtype=torch.half,
+                    device="cuda").normal_().requires_grad_())
+
+    head_kv = H // groups
+    K = (
+        torch.empty(BATCH, N_CTX, head_kv, D_HEAD_QK, dtype=torch.half,
+                    device="cuda").normal_().requires_grad_())
+    V = (
+        torch.empty(BATCH, N_CTX, head_kv, D_HEAD_V, dtype=torch.half,
+                    device="cuda").normal_().requires_grad_())
+    dO = (
+        torch.empty(BATCH, N_CTX, H, D_HEAD_V, dtype=torch.half,
+                    device="cuda").normal_().requires_grad_())
+
+    M_BLOCKS = N_CTX // BLOCK_M
+    N_BLOCKS = N_CTX // BLOCK_N
+
+    block_mask = torch.randint(
+        low=0, high=2, size=(BATCH, H, M_BLOCKS, N_BLOCKS), dtype=torch.bool, device="cuda")
+
+    attention = BlockSparseAttentionKernel(
+        BATCH, H, N_CTX, D_HEAD_QK, D_HEAD_V, BLOCK_M, BLOCK_N, causal=causal, groups=groups)
+
+    o = attention.backward(Q, K, V, dO, block_mask)
+    print(o)
+
+    attention.check(Q, K, V, dO, block_mask)
+    attention.profile(Q, K, V, dO, block_mask)
+
+
+if __name__ == "__main__":
+    main()

tileops-0.0.1.dev0/tests/test_gqa.py

@@ -0,0 +1,49 @@
+# Copyright (c) Tile-AI.
+# Licensed under the MIT License.
+
+import argparse
+import torch
+from top import GQAKernel
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=8, help='batch size')
+    parser.add_argument('--h', type=int, default=32, help='Number of heads')
+    parser.add_argument('--n_ctx', type=int, default=1024, help='Context size')
+    parser.add_argument('--d_head_qk', type=int, default=128, help='Head dimension for Q/K')
+    parser.add_argument('--d_head_v', type=int, default=64, help='Head dimension for V')
+    parser.add_argument('--causal', type=bool, default=False, help='Casual flag')
+    parser.add_argument('--groups', type=int, default=16, help='groups')
+    args = parser.parse_args()
+    BATCH, H, N_CTX, D_HEAD_QK, D_HEAD_V, groups, causal = args.batch, args.h, args.n_ctx, args.d_head_qk, args.d_head_v, args.groups, args.causal
+
+    BLOCK_M = 128
+    BLOCK_N = 64
+    # q = torch.randn(batch, heads, dim, device='cuda', dtype=torch.float16)
+    Q = (
+        torch.empty(BATCH, N_CTX, H, D_HEAD_QK, dtype=torch.half,
+                    device="cuda").normal_().requires_grad_())
+
+    head_kv = H // groups
+    K = (
+        torch.empty(BATCH, N_CTX, head_kv, D_HEAD_QK, dtype=torch.half,
+                    device="cuda").normal_().requires_grad_())
+    V = (
+        torch.empty(BATCH, N_CTX, head_kv, D_HEAD_V, dtype=torch.half,
+                    device="cuda").normal_().requires_grad_())
+    dO = (
+        torch.empty(BATCH, N_CTX, H, D_HEAD_V, dtype=torch.half,
+                    device="cuda").normal_().requires_grad_())
+
+    gqa = GQAKernel(
+        BATCH, H, N_CTX, D_HEAD_QK, D_HEAD_V, BLOCK_M, BLOCK_N, causal=causal, groups=groups)
+    o = gqa.backward(Q, K, V, dO)
+    print(o)
+    latency = gqa.profile()
+    print(f"Latency: {latency:.4f} ms")
+    gqa.check(Q, K, V, dO)
+
+
+if __name__ == "__main__":
+    main()

tileops-0.0.1.dev0/tests/test_gqa_decode.py

@@ -0,0 +1,42 @@
+import argparse
+import torch
+from top import GQADecodeKernel
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=1, help='batch size')
+    parser.add_argument('--heads', type=int, default=32, help='heads')
+    parser.add_argument('--groups', type=int, default=8, help='groups')
+    parser.add_argument('--kv_seqlen', type=int, default=8192, help='kv sequence length')
+    parser.add_argument('--dim', type=int, default=128, help='dim')
+    parser.add_argument('--num_split', type=int, default=8, help='num_split')
+    parser.add_argument('--tune', action='store_true', help='tune configs')
+    args = parser.parse_args()
+    batch, heads, groups, kv_seqlen, dim, num_split, tune = args.batch, args.heads, args.groups, args.kv_seqlen, args.dim, args.num_split, args.tune
+
+    BLOCK_N = 128
+    BLOCK_H = 64
+    threads = 128
+
+    Q = torch.randn(batch, heads, dim, device="cuda", dtype=torch.float16)
+    K = torch.randn(batch, kv_seqlen, groups, dim, device="cuda", dtype=torch.float16)
+    V = torch.randn(batch, kv_seqlen, groups, dim, device="cuda", dtype=torch.float16)
+
+    gqa_decode = GQADecodeKernel(batch, heads, kv_seqlen, dim, BLOCK_N, BLOCK_H, threads, num_split,
+                                 groups)
+    if tune:
+        gqa_decode.autotune()
+    o = gqa_decode.decode(Q, K, V)
+    print(o)
+
+    gqa_decode.profile()
+
+    gqa_decode.check(Q, K, V)
+
+    gqa_decode.autotune()
+    gqa_decode.profile()
+
+
+if __name__ == "__main__":
+    main()

tileops-0.0.1.dev0/tests/test_linear_attn_fused_chunk.py

@@ -0,0 +1,23 @@
+# Copyright (c) Tile-AI.
+# Licensed under the MIT License.
+
+import argparse
+from top import LinearAttentionFusedChunkKernel
+
+
+def test_fused_chunk(B, S, H, D):
+    kernel = LinearAttentionFusedChunkKernel(B, S, H, D)
+    kernel.check()
+    kernel.profile()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=8, help='batch size')
+    parser.add_argument('--seq_len', type=int, default=2048, help='sequence length')
+    parser.add_argument('--heads', type=int, default=32, help='num heads')
+    parser.add_argument('--dim', type=int, default=256, help='head dim')
+    args = parser.parse_args()
+    B, S, H, D = args.batch, args.seq_len, args.heads, args.dim
+
+    test_fused_chunk(B, S, H, D)

tileops-0.0.1.dev0/tests/test_linear_attn_fused_recurrent.py

@@ -0,0 +1,23 @@
+# Copyright (c) Tile-AI.
+# Licensed under the MIT License.
+
+import argparse
+from top import LinearAttentionFusedRecurrentKernel
+
+
+def test_fused_recurrent(B, S, H, D):
+    kernel = LinearAttentionFusedRecurrentKernel(B, S, H, D)
+    kernel.check()
+    kernel.profile()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=8, help='batch size')
+    parser.add_argument('--seq_len', type=int, default=2048, help='sequence length')
+    parser.add_argument('--heads', type=int, default=32, help='num heads')
+    parser.add_argument('--dim', type=int, default=256, help='head dim')
+    args = parser.parse_args()
+    B, S, H, D = args.batch, args.seq_len, args.heads, args.dim
+
+    test_fused_recurrent(B, S, H, D)

tileops-0.0.1.dev0/tests/test_mamba_chunk_scan.py

@@ -0,0 +1,71 @@
+# Copyright (c) Tile-AI.
+# Licensed under the MIT License.
+
+import argparse
+import torch
+from top import MambaChunkScanKernel
+import math
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=8, help='batch size')
+    parser.add_argument('--heads', type=int, default=80, help='heads')
+    parser.add_argument('--groups', type=int, default=1, help='groups')
+    parser.add_argument('--seq_len', type=int, default=4096, help='sequence length')
+    parser.add_argument('--chunk_size', type=int, default=256, help='chunk size')
+    parser.add_argument('--dim', type=int, default=64, help='dim')
+    parser.add_argument('--dstate', type=int, default=128, help='dstate')
+    parser.add_argument('--tune', action='store_true', help='tune configs')
+    args = parser.parse_args()
+    BATCH, HEADS, GROUPS, SEQ_LEN, CHUNK_SIZE, DIM, DSTATE = args.batch, args.heads, args.groups, args.seq_len, args.chunk_size, args.dim, args.dstate
+    NCHUNKS = math.ceil(SEQ_LEN / CHUNK_SIZE)
+
+    cb = torch.empty((BATCH, NCHUNKS, GROUPS, CHUNK_SIZE, CHUNK_SIZE),
+                     dtype=torch.half).cuda().normal_(-1.0, 1.0)
+    x = torch.empty((BATCH, SEQ_LEN, HEADS, DIM), dtype=torch.half).cuda().normal_(-1.0, 1.0)
+    dt = torch.empty((BATCH, HEADS, NCHUNKS, CHUNK_SIZE),
+                     dtype=torch.half).cuda().normal_(-1.0, 1.0)
+    dA_cumsum = torch.empty((BATCH, HEADS, NCHUNKS, CHUNK_SIZE),
+                            dtype=torch.half).cuda().normal_(-1.0, 1.0)
+    C = torch.empty((BATCH, SEQ_LEN, GROUPS, DSTATE), dtype=torch.half).cuda().normal_(-1.0, 1.0)
+    prev_states = torch.empty((BATCH, NCHUNKS, HEADS, DIM, DSTATE),
+                              dtype=torch.half).cuda().normal_(-1.0, 1.0)
+    D = torch.empty((HEADS,), dtype=torch.half).cuda().normal_(-1.0, 1.0)
+
+    if args.tune:
+        mamba_chunk_scan = MambaChunkScanKernel(
+            BATCH, HEADS, GROUPS, SEQ_LEN, CHUNK_SIZE, DIM, DSTATE, tune=True)
+
+    else:
+        block_M = 64
+        block_N = 64
+        block_K = 64
+        block_Dstate = 128
+        num_stages = 2
+        threads = 128
+
+        mamba_chunk_scan = MambaChunkScanKernel(
+            BATCH,
+            HEADS,
+            GROUPS,
+            SEQ_LEN,
+            CHUNK_SIZE,
+            DIM,
+            DSTATE,
+            block_M=block_M,
+            block_N=block_N,
+            block_K=block_K,
+            block_Dstate=block_Dstate,
+            num_stages=num_stages,
+            threads=threads)
+
+    o = mamba_chunk_scan(cb, x, dt, dA_cumsum, C, prev_states, D)
+    print(o)
+    latency = mamba_chunk_scan.profile()
+    print(f"Latency: {latency:.4f} ms")
+    mamba_chunk_scan.check(cb, x, dt, dA_cumsum, C, prev_states, D)
+
+
+if __name__ == "__main__":
+    main()

tileops-0.0.1.dev0/tests/test_mamba_chunk_state.py

@@ -0,0 +1,64 @@
+# Copyright (c) Tile-AI.
+# Licensed under the MIT License.
+
+import argparse
+import torch
+from top import MambaChunkStateKernel
+import math
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=8, help='batch size')
+    parser.add_argument('--heads', type=int, default=80, help='heads')
+    parser.add_argument('--groups', type=int, default=1, help='groups')
+    parser.add_argument('--seq_len', type=int, default=4096, help='sequence length')
+    parser.add_argument('--chunk_size', type=int, default=256, help='chunk size')
+    parser.add_argument('--dim', type=int, default=64, help='dim')
+    parser.add_argument('--dstate', type=int, default=128, help='dstate')
+    parser.add_argument('--tune', action='store_true', help='tune configs')
+    args = parser.parse_args()
+    BATCH, HEADS, GROUPS, SEQ_LEN, CHUNK_SIZE, DIM, DSTATE = args.batch, args.heads, args.groups, args.seq_len, args.chunk_size, args.dim, args.dstate
+    NCHUNKS = math.ceil(SEQ_LEN / CHUNK_SIZE)
+
+    B = torch.empty((BATCH, SEQ_LEN, GROUPS, DSTATE), dtype=torch.half).cuda().normal_(-1.0, 1.0)
+    x = torch.empty((BATCH, SEQ_LEN, HEADS, DIM), dtype=torch.half).cuda().normal_(-1.0, 1.0)
+    dt = torch.empty((BATCH, HEADS, NCHUNKS, CHUNK_SIZE),
+                     dtype=torch.half).cuda().normal_(-1.0, 1.0)
+    dA_cumsum = torch.empty((BATCH, HEADS, NCHUNKS, CHUNK_SIZE),
+                            dtype=torch.half).cuda().normal_(-1.0, 1.0)
+
+    if args.tune:
+        mamba_chunk_state = MambaChunkStateKernel(
+            BATCH, HEADS, GROUPS, SEQ_LEN, CHUNK_SIZE, DIM, DSTATE, tune=True)
+
+    else:
+        block_M = 64
+        block_N = 128
+        block_K = 64
+        num_stages = 4
+        threads = 128
+
+        mamba_chunk_state = MambaChunkStateKernel(
+            BATCH,
+            HEADS,
+            GROUPS,
+            SEQ_LEN,
+            CHUNK_SIZE,
+            DIM,
+            DSTATE,
+            block_M=block_M,
+            block_N=block_N,
+            block_K=block_K,
+            num_stages=num_stages,
+            threads=threads)
+
+    o = mamba_chunk_state(B, x, dt, dA_cumsum)
+    print(o)
+    latency = mamba_chunk_state.profile()
+    print(f"Latency: {latency:.4f} ms")
+    mamba_chunk_state.check(B, x, dt, dA_cumsum)
+
+
+if __name__ == "__main__":
+    main()

tileops-0.0.1.dev0/tests/test_mha.py

@@ -0,0 +1,24 @@
+# Copyright (c) Tile-AI.
+# Licensed under the MIT License.
+
+import argparse
+from top import MHAKernel
+
+
+def test_mha_kernel(B, S, H, D, causal):
+    kernel = MHAKernel(B, H, S, D, causal)
+    kernel.check()
+    kernel.profile()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=8, help='batch size')
+    parser.add_argument('--seq_len', type=int, default=1024, help='sequence length')
+    parser.add_argument('--heads', type=int, default=32, help='num heads')
+    parser.add_argument('--dim', type=int, default=64, help='head dim')
+    parser.add_argument('--causal', action='store_true', default=False, help='causal attention')
+    args = parser.parse_args()
+    B, S, H, D, causal = args.batch, args.seq_len, args.heads, args.dim, args.causal
+
+    test_mha_kernel(B, S, H, D, causal)

tileops-0.0.1.dev0/tests/test_mha_decode.py

@@ -0,0 +1,23 @@
+# Copyright (c) Tile-AI.
+# Licensed under the MIT License.
+
+import argparse
+from top import MHADecodeKernel
+
+
+def test_mha_decode_kernel(B, S, H, D):
+    kernel = MHADecodeKernel(B, H, S, D)
+    kernel.check()
+    kernel.profile()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=1, help='batch size')
+    parser.add_argument('--seqlen_kv', type=int, default=8192, help='sequence length')
+    parser.add_argument('--heads', type=int, default=32, help='num heads')
+    parser.add_argument('--dim', type=int, default=128, help='head dim')
+    args = parser.parse_args()
+    B, S, H, D = args.batch, args.seqlen_kv, args.heads, args.dim
+
+    test_mha_decode_kernel(B, S, H, D)

tileops-0.0.1.dev0/tests/test_mla.py

@@ -0,0 +1,39 @@
+# Copyright (c) Tile-AI.
+# Licensed under the MIT License.
+
+import argparse
+import torch
+from top import MLAKernel
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=128, help='batch size')
+    parser.add_argument('--heads', type=int, default=128, help='q heads number')
+    parser.add_argument('--kv_heads', type=int, default=1, help='kv heads number')
+    parser.add_argument('--kv_ctx', type=int, default=4096, help='kv context length')
+    parser.add_argument('--dim', type=int, default=256, help='head dim')
+    parser.add_argument('--pe_dim', type=int, default=64, help='pe head dim')
+    args = parser.parse_args()
+    batch, heads, kv_heads, kv_ctx, dim, pe_dim = args.batch, args.heads, args.kv_heads, args.kv_ctx, args.dim, args.pe_dim
+
+    BLOCK_N = 32
+    BLOCK_H = 32
+    num_split = 1
+
+    mla = MLAKernel(batch, heads, kv_heads, kv_ctx, dim, pe_dim, BLOCK_N, BLOCK_H, num_split)
+
+    q = torch.randn(batch, heads, dim, device='cuda', dtype=torch.float16)
+    q_pe = torch.randn(batch, heads, pe_dim, device='cuda', dtype=torch.float16)
+    kv = torch.randn(batch, kv_ctx, kv_heads, dim, device='cuda', dtype=torch.float16)
+    k_pe = torch.randn(batch, kv_ctx, kv_heads, pe_dim, device='cuda', dtype=torch.float16)
+
+    o = mla(q, q_pe, kv, k_pe)
+    print(o)
+    latency = mla.profile()
+    print(f"Latency: {latency:.4f} ms")
+    mla.check(q, q_pe, kv, k_pe)
+
+
+if __name__ == "__main__":
+    main()

tileops-0.0.1.dev0/tileops.egg-info/PKG-INFO

@@ -0,0 +1,82 @@
+Metadata-Version: 2.4
+Name: tileops
+Version: 0.0.1.dev0
+Summary: TileOPs kernels for efficient inference
+Author: Tile-AI
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/tile-ai/tileops
+Project-URL: Bug Tracker, https://github.com/tile-ai/tileops/issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: tilelang
+Dynamic: license-file
+
+# TileOPs (TOP)
+
+**TileOPs (TOP)** is a high-performance machine learning operator collections built to run on the [TileLang](https://github.com/tile-ai/tilelang) backend. It offers efficient, modular, and composable implementations optimized for AI workloads.
+
+---
+
+
+## 📦 Installation
+
+### Requirements
+
+- Python 3.8+
+- PyTorch >= 2.1
+- [TileLang](https://github.com/tilelang/tilelang)
+- Triton (optional, for selected fast kernels)
+
+### Install (editable mode for development)
+
+```bash
+git clone https://github.com/tile-ai/TileOPs
+cd TileOPs
+git submodule update --init --recursive
+pip install -e .
+```
+
+## 🚀 Quick Usage
+
+```python
+import torch
+import top
+from top import MLAKernel
+
+device = "cuda"
+dtype = torch.float16
+
+batch = 128
+heads = 64
+kv_heads = 1
+kv_ctx = 8192
+dim = 512
+pe_dim = 64
+
+# Query input: [batch, heads, dim]
+q = torch.randn(batch, heads, dim, device=device, dtype=dtype)
+
+# Query positional encoding: [batch, heads, pe_dim]
+q_pe = torch.randn(batch, heads, pe_dim, device=device, dtype=dtype)
+
+# KV cache input: [batch, kv_ctx, kv_heads, dim]
+kv = torch.randn(batch, kv_ctx, kv_heads, dim, device=device, dtype=dtype)
+
+# KV positional encoding: [batch, kv_ctx, kv_heads, pe_dim]
+k_pe = torch.randn(batch, kv_ctx, kv_heads, pe_dim, device=device, dtype=dtype)
+
+# Use MLA kernel
+block_N = 64
+block_H = 64
+num_split = 1
+
+mla = MLAKernel(batch, heads, kv_heads, kv_ctx, dim, pe_dim, block_N, block_H, num_split)
+
+out = mla(q, q_pe, kv, k_pe)
+```
+
+## Acknowledgments

tileops-0.0.1.dev0/tileops.egg-info/SOURCES.txt

@@ -0,0 +1,21 @@
+LICENSE
+README.md
+pyproject.toml
+setup.py
+tests/test_bitnet.py
+tests/test_blocksparse_attention.py
+tests/test_gqa.py
+tests/test_gqa_decode.py
+tests/test_linear_attn_fused_chunk.py
+tests/test_linear_attn_fused_recurrent.py
+tests/test_mamba_chunk_scan.py
+tests/test_mamba_chunk_state.py
+tests/test_mha.py
+tests/test_mha_decode.py
+tests/test_mla.py
+tileops.egg-info/PKG-INFO
+tileops.egg-info/SOURCES.txt
+tileops.egg-info/dependency_links.txt
+tileops.egg-info/requires.txt
+tileops.egg-info/top_level.txt
+top/__init__.py

tileops-0.0.1.dev0/tileops.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

tileops-0.0.1.dev0/tileops.egg-info/requires.txt

@@ -0,0 +1,2 @@
+torch
+tilelang

tileops-0.0.1.dev0/tileops.egg-info/top_level.txt

@@ -0,0 +1 @@
+top

tileops-0.0.1.dev0/top/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) Tile-AI.
+# Licensed under the MIT License.
+
+from .kernel.mha import MHAKernel, MHADecodeKernel
+from .kernel.mla import MLAKernel
+from .kernel.gqa import GQAKernel, GQADecodeKernel
+from .kernel.mamba_chunk_scan import MambaChunkScanKernel
+from .kernel.mamba_chunk_state import MambaChunkStateKernel
+from .kernel.blocksparse_attention import BlockSparseAttentionKernel
+from .kernel.linear_attention.linear_attn import LinearAttentionFusedChunkKernel, LinearAttentionFusedRecurrentKernel
+from .kernel.bitnet import Bitnet_158_int8xint2_kernel
+
+__all__ = [
+    "MHAKernel", "MHADecodeKernel", "MLAKernel", "GQAKernel", "GQADecodeKernel",
+    "MambaChunkScanKernel", "MambaChunkStateKernel", "BlockSparseAttentionKernel",
+    "LinearAttentionFusedChunkKernel", "LinearAttentionFusedRecurrentKernel",
+    "Bitnet_158_int8xint2_kernel"
+]