tileops 0.0.1.dev0__tar.gz → 0.0.1.dev1__tar.gz

tileops-0.0.1.dev1/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.4
+Name: tileops
+Version: 0.0.1.dev1
+Summary: TileOPs kernels for efficient inference
+Author: Tile-AI
+License-Expression: MIT
+Requires-Python: >=3.8
+License-File: LICENSE
+Requires-Dist: torch>=2.1.0
+Requires-Dist: numpy>=1.23.5
+Requires-Dist: einops>=0.6.0
+Requires-Dist: Cython
+Requires-Dist: tqdm>=4.62.3
+Requires-Dist: typing_extensions>=4.10.0
+Requires-Dist: cloudpickle
+Requires-Dist: ml_dtypes
+Requires-Dist: psutil
+Requires-Dist: tilelang>=0.1.6.post1
+Provides-Extra: dev
+Requires-Dist: yapf==0.40.2; extra == "dev"
+Requires-Dist: toml==0.10.2; extra == "dev"
+Requires-Dist: tomli==2.0.1; extra == "dev"
+Requires-Dist: ruff==0.6.5; extra == "dev"
+Requires-Dist: codespell==2.3.0; extra == "dev"
+Requires-Dist: clang-format==15.0.7; extra == "dev"
+Dynamic: license-file

tileops-0.0.1.dev1/README.md

@@ -0,0 +1,119 @@
+# TileOPs (TOP)
+
+**TileOPs (TOP)** is a high-performance machine learning operator collections built on top of [TileLang](https://github.com/tile-ai/tilelang). It offers efficient, modular, and composable implementations optimized for AI workloads.
+
+Note: TileOPs is still under rapid development.
+
+---
+
+![DeepSeek-V3.2-Exp DeepSeek Sparse Attention (DSA) performance on H800 SXM](https://raw.githubusercontent.com/tile-ai/TileOPs/main/docs/figures/sparse_mla_perf.png)
+*DeepSeek-V3.2-Exp DeepSeek Sparse Attention (DSA) performance on H800 SXM*
+
+## 📦 Installation
+
+### Requirements
+
+- Python 3.8+
+- PyTorch >= 2.1
+- GLIBCXX_3.4.32
+- [TileLang](https://github.com/tile-ai/tilelang)
+
+### Method 1: Install with Pip
+
+```bash
+pip install tileops
+```
+
+### Method 2: Install from source (editable mode for development)
+
+```bash
+git clone https://github.com/tile-ai/TileOPs
+cd TileOPs
+pip install -e '.[dev]' -v # remove -e option if you don't want to install in editable mode, -v for verbose output
+```
+
+## 🚀 Quick Usage
+
+### Sparse MLA
+
+```python
+import torch
+from top import SparseMLAKernel
+
+batch_size = 1
+seq_len = 1024
+seq_len_kv = 2048
+q_start_index_s = 1024
+n_heads = 128
+head_dim = 512
+tail_dim = 64
+topk = 2048
+kv_stride = 1
+kv_group = 1
+sm_scale = None
+
+sparse_mla = SparseMLAKernel(
+    batch=batch_size,
+    seq_len=seq_len,
+    seq_len_kv=seq_len_kv,
+    q_start_index_s=q_start_index_s,
+    heads=n_heads,
+    dim=head_dim,
+    tail_dim=tail_dim,
+    topk=topk,
+    kv_stride=kv_stride,
+    kv_group=kv_group,
+    sm_scale=sm_scale,
+    is_casual=True,
+    dtype=torch.bfloat16,
+    device='cuda',
+)
+
+# Evaluate the Sparse MLA kernel performance
+sparse_mla.check()
+latency = sparse_mla.profile()
+print(f"Latency: {latency:.4f} ms")
+print(f'fwd tflops = ',
+        (batch_size * seq_len * (head_dim + tail_dim + head_dim) * topk * 2 * n_heads) / (latency * 1e-3) / 1e12)
+```
+
+### MLA
+
+```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-0.0.1.dev1}/pyproject.toml

@@ -1,21 +1,41 @@
 [project]
 name = "tileops"
-version = "0.0.1.dev0"
+version = "0.0.1.dev1"
 description = "TileOPs kernels for efficient inference"
-authors = [{ name="Tile-AI" }]
-readme = "README.md"
+authors = [{ name="Tile-AI"}]
+dependencies = [    
+    "torch>=2.1.0",
+    "numpy>=1.23.5",
+    "einops>=0.6.0",
+    "Cython",
+    "tqdm>=4.62.3",
+    "typing_extensions>=4.10.0",
+    "cloudpickle",
+    "ml_dtypes",
+    "psutil",
+    "tilelang>=0.1.6.post1"
+]
 requires-python = ">=3.8"
 license = "MIT"
-dependencies = ["torch", "tilelang"]
-classifiers = [
-    "Development Status :: 3 - Alpha",
-    "Programming Language :: Python :: 3"
+
+[project.optional-dependencies]
+dev = [
+"yapf==0.40.2",
+"toml==0.10.2",
+"tomli==2.0.1",
+"ruff==0.6.5",
+"codespell==2.3.0",
+"clang-format==15.0.7",
 ]
 
 [build-system]
-requires = ["setuptools", "wheel"]
+requires = ["setuptools>=68", "wheel", "build"]
 build-backend = "setuptools.build_meta"
 
+[tool.setuptools.packages.find] 
+where = ["."] 
+exclude = ["3rdparty", "3rdparty.*", "tests", "tests.*"]
+
 [tool.yapf]
 based_on_style = "yapf"
 column_limit = 100
@@ -30,13 +50,6 @@ skip = [
     ".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

tileops-0.0.1.dev1/setup.py

@@ -0,0 +1,44 @@
+import os
+import sys
+from setuptools import setup
+
+
+# Control whether to use custom installation logic
+# Custom installation is enabled only when:
+# 1. TILEOPS_DEV_INSTALL environment variable is set to "1"
+# 2. The current command is not for building distribution packages (bdist_wheel, sdist)
+# This ensures that custom installation logic is skipped during package building
+USE_CUSTOM_INSTALL = (
+os.environ.get("TILEOPS_DEV_INSTALL") == "1"
+and not any(cmd in sys.argv for cmd in ["bdist_wheel", "sdist"])
+)
+
+cmdclass = {}
+
+if USE_CUSTOM_INSTALL:
+    from setuptools.command.install import install
+    import subprocess
+    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)
+
+    cmdclass = {"install": CustomInstall}
+
+setup(name = "tileops",
+      version = "0.0.1.dev1",
+      cmdclass=cmdclass)

{tileops-0.0.1.dev0 → tileops-0.0.1.dev1}/tests/test_blocksparse_attention.py

@@ -1,6 +1,3 @@
-# Copyright (c) Tile-AI.
-# Licensed under the MIT License.
-
 import argparse
 import torch
 from top import BlockSparseAttentionKernel
@@ -42,13 +39,23 @@ def main():
         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)
+        BATCH,
+        H,
+        N_CTX,
+        D_HEAD_QK,
+        D_HEAD_V,
+        BLOCK_M,
+        BLOCK_N,
+        causal=causal,
+        block_mask=block_mask,
+        groups=groups,
+        fwd_tune=True,
+        bwd_tune=True)
+
+    o = attention.backward(Q, K, V, dO)
     print(o)
 
-    attention.check(Q, K, V, dO, block_mask)
-    attention.profile(Q, K, V, dO, block_mask)
+    attention.profile()
 
 
 if __name__ == "__main__":

{tileops-0.0.1.dev0 → tileops-0.0.1.dev1}/tests/test_gqa.py

@@ -1,6 +1,3 @@
-# Copyright (c) Tile-AI.
-# Licensed under the MIT License.
-
 import argparse
 import torch
 from top import GQAKernel
@@ -18,8 +15,10 @@ def main():
     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
+    fwd_BLOCK_M = 128
+    fwd_BLOCK_N = 32
+    bwd_BLOCK_M = 64
+    bwd_BLOCK_N = 32
     # q = torch.randn(batch, heads, dim, device='cuda', dtype=torch.float16)
     Q = (
         torch.empty(BATCH, N_CTX, H, D_HEAD_QK, dtype=torch.half,
@@ -37,11 +36,20 @@ def main():
                     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")
+        BATCH,
+        H,
+        N_CTX,
+        D_HEAD_QK,
+        D_HEAD_V,
+        fwd_BLOCK_M,
+        fwd_BLOCK_N,
+        bwd_BLOCK_M,
+        bwd_BLOCK_N,
+        fwd_tune=True,
+        bwd_tune=True,
+        causal=causal,
+        groups=groups)
+    gqa.profile(Q, K, V, dO)
     gqa.check(Q, K, V, dO)
 
 

{tileops-0.0.1.dev0 → tileops-0.0.1.dev1}/tests/test_linear_attn_fused_chunk.py

@@ -1,6 +1,3 @@
-# Copyright (c) Tile-AI.
-# Licensed under the MIT License.
-
 import argparse
 from top import LinearAttentionFusedChunkKernel
 

{tileops-0.0.1.dev0 → tileops-0.0.1.dev1}/tests/test_linear_attn_fused_recurrent.py

@@ -1,6 +1,3 @@
-# Copyright (c) Tile-AI.
-# Licensed under the MIT License.
-
 import argparse
 from top import LinearAttentionFusedRecurrentKernel
 

{tileops-0.0.1.dev0 → tileops-0.0.1.dev1}/tests/test_mamba_chunk_scan.py

@@ -1,6 +1,3 @@
-# Copyright (c) Tile-AI.
-# Licensed under the MIT License.
-
 import argparse
 import torch
 from top import MambaChunkScanKernel

{tileops-0.0.1.dev0 → tileops-0.0.1.dev1}/tests/test_mamba_chunk_state.py

@@ -1,6 +1,3 @@
-# Copyright (c) Tile-AI.
-# Licensed under the MIT License.
-
 import argparse
 import torch
 from top import MambaChunkStateKernel

{tileops-0.0.1.dev0 → tileops-0.0.1.dev1}/tests/test_mha.py

@@ -1,12 +1,9 @@
-# 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)
+def test_mha_kernel(B, S, H, D, causal, tune):
+    kernel = MHAKernel(B, H, S, D, causal, fwd_tune=tune, bwd_tune=tune)
     kernel.check()
     kernel.profile()
 
@@ -18,7 +15,8 @@ if __name__ == "__main__":
     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')
+    parser.add_argument('--tune', action='store_true', default=True, help='tune the kernel')
     args = parser.parse_args()
-    B, S, H, D, causal = args.batch, args.seq_len, args.heads, args.dim, args.causal
+    B, S, H, D, causal, tune = args.batch, args.seq_len, args.heads, args.dim, args.causal, args.tune
 
-    test_mha_kernel(B, S, H, D, causal)
+    test_mha_kernel(B, S, H, D, causal, tune)

{tileops-0.0.1.dev0 → tileops-0.0.1.dev1}/tests/test_mha_decode.py

@@ -1,12 +1,9 @@
-# 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)
+def test_mha_decode_kernel(B, S, H, D, tune):
+    kernel = MHADecodeKernel(B, H, S, D, tune=tune)
     kernel.check()
     kernel.profile()
 
@@ -17,7 +14,8 @@ if __name__ == "__main__":
     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')
+    parser.add_argument('--tune', action='store_true', default=True, help='tune the kernel')
     args = parser.parse_args()
-    B, S, H, D = args.batch, args.seqlen_kv, args.heads, args.dim
+    B, S, H, D, tune = args.batch, args.seqlen_kv, args.heads, args.dim, args.tune
 
-    test_mha_decode_kernel(B, S, H, D)
+    test_mha_decode_kernel(B, S, H, D, tune)

{tileops-0.0.1.dev0 → tileops-0.0.1.dev1}/tests/test_mla.py

@@ -1,6 +1,3 @@
-# Copyright (c) Tile-AI.
-# Licensed under the MIT License.
-
 import argparse
 import torch
 from top import MLAKernel

tileops-0.0.1.dev1/tests/test_sparse_mla.py

@@ -0,0 +1,57 @@
+import argparse
+import torch
+from top import SparseMLAKernel
+
+
+def test_sparsemla_kernel(batch, seq_len, seq_len_kv, q_start_index_s, heads, dim, tail_dim, topk,
+                          kv_stride, kv_group, sm_scale, dtype):
+    sparse_mla = SparseMLAKernel(
+        batch=batch,
+        seq_len=seq_len,
+        seq_len_kv=seq_len_kv,
+        q_start_index_s=q_start_index_s,
+        heads=heads,
+        dim=dim,
+        tail_dim=tail_dim,
+        topk=topk,
+        kv_stride=kv_stride,
+        kv_group=kv_group,
+        sm_scale=sm_scale,
+        is_casual=True,
+        dtype=dtype,
+        device='cuda',
+    )
+    sparse_mla.check()
+    latency = sparse_mla.profile()
+    print(f"Latency: {latency:.4f} ms")
+    print(f'fwd tflops = ',
+          (batch * seq_len * (dim + tail_dim + dim) * topk * 2 * heads) / (latency * 1e-3) / 1e12)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=1, help='batch size')
+    parser.add_argument('--seq_len', type=int, default=1024, help='sequence length')
+    parser.add_argument('--seq_len_kv', type=int, default=2048, help='key/value sequence length')
+    parser.add_argument('--heads', type=int, default=128, help='num heads')
+    parser.add_argument('--dim', type=int, default=512, help='head dim')
+    parser.add_argument('--tail_dim', type=int, default=64, help='tail dim')
+    parser.add_argument('--topk', type=int, default=2048, help='topk')
+    parser.add_argument('--kv_stride', type=int, default=1, help='kv_stride')
+    parser.add_argument('--kv_group', type=int, default=1, help='kv_group')
+    parser.add_argument('--sm_scale', type=float, default=None, help='softmax scaling factor')
+    args = parser.parse_args()
+    batch = args.batch
+    seq_len = args.seq_len
+    seq_len_kv = args.seq_len_kv
+    heads = args.heads
+    dim = args.dim
+    tail_dim = args.tail_dim
+    topk = args.topk
+    kv_stride = args.kv_stride
+    kv_group = args.kv_group
+    sm_scale = args.sm_scale
+    dtype = torch.bfloat16
+
+    test_sparsemla_kernel(batch, seq_len, seq_len_kv, 1024, heads, dim, tail_dim, topk, kv_stride,
+                          kv_group, sm_scale, dtype)

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

@@ -0,0 +1,26 @@
+Metadata-Version: 2.4
+Name: tileops
+Version: 0.0.1.dev1
+Summary: TileOPs kernels for efficient inference
+Author: Tile-AI
+License-Expression: MIT
+Requires-Python: >=3.8
+License-File: LICENSE
+Requires-Dist: torch>=2.1.0
+Requires-Dist: numpy>=1.23.5
+Requires-Dist: einops>=0.6.0
+Requires-Dist: Cython
+Requires-Dist: tqdm>=4.62.3
+Requires-Dist: typing_extensions>=4.10.0
+Requires-Dist: cloudpickle
+Requires-Dist: ml_dtypes
+Requires-Dist: psutil
+Requires-Dist: tilelang>=0.1.6.post1
+Provides-Extra: dev
+Requires-Dist: yapf==0.40.2; extra == "dev"
+Requires-Dist: toml==0.10.2; extra == "dev"
+Requires-Dist: tomli==2.0.1; extra == "dev"
+Requires-Dist: ruff==0.6.5; extra == "dev"
+Requires-Dist: codespell==2.3.0; extra == "dev"
+Requires-Dist: clang-format==15.0.7; extra == "dev"
+Dynamic: license-file

{tileops-0.0.1.dev0 → tileops-0.0.1.dev1}/tileops.egg-info/SOURCES.txt

@@ -13,9 +13,23 @@ tests/test_mamba_chunk_state.py
 tests/test_mha.py
 tests/test_mha_decode.py
 tests/test_mla.py
+tests/test_sparse_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
+top/__init__.py
+top/kernel/__init__.py
+top/kernel/bitnet.py
+top/kernel/blocksparse_attention.py
+top/kernel/gqa.py
+top/kernel/mamba_chunk_scan.py
+top/kernel/mamba_chunk_state.py
+top/kernel/mha.py
+top/kernel/mla.py
+top/kernel/sparse_mla.py
+top/kernel/linear_attention/__init__.py
+top/kernel/linear_attention/linear_attn.py
+top/utils/__init__.py
+top/utils/utils.py

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

@@ -0,0 +1,18 @@
+torch>=2.1.0
+numpy>=1.23.5
+einops>=0.6.0
+Cython
+tqdm>=4.62.3
+typing_extensions>=4.10.0
+cloudpickle
+ml_dtypes
+psutil
+tilelang>=0.1.6.post1
+
+[dev]
+yapf==0.40.2
+toml==0.10.2
+tomli==2.0.1
+ruff==0.6.5
+codespell==2.3.0
+clang-format==15.0.7

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

@@ -0,0 +1,4 @@
+build
+dist
+docs
+top

{tileops-0.0.1.dev0 → tileops-0.0.1.dev1}/top/__init__.py

@@ -1,18 +1,16 @@
-# 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.linear_attention.linear_attn import LinearAttentionFusedChunkKernel, LinearAttentionFusedRecurrentKernel
 from .kernel.bitnet import Bitnet_158_int8xint2_kernel
+from .kernel.sparse_mla import SparseMLAKernel
 
 __all__ = [
     "MHAKernel", "MHADecodeKernel", "MLAKernel", "GQAKernel", "GQADecodeKernel",
     "MambaChunkScanKernel", "MambaChunkStateKernel", "BlockSparseAttentionKernel",
-    "LinearAttentionFusedChunkKernel", "LinearAttentionFusedRecurrentKernel",
-    "Bitnet_158_int8xint2_kernel"
+    # "LinearAttentionFusedChunkKernel", "LinearAttentionFusedRecurrentKernel",
+    "Bitnet_158_int8xint2_kernel", "SparseMLAKernel"
 ]

tileops-0.0.1.dev1/top/kernel/__init__.py

@@ -0,0 +1,15 @@
+from .mha import MHAKernel, MHADecodeKernel
+from .mla import MLAKernel
+from .gqa import GQAKernel, GQADecodeKernel
+from .mamba_chunk_scan import MambaChunkScanKernel
+from .mamba_chunk_state import MambaChunkStateKernel
+from .blocksparse_attention import BlockSparseAttentionKernel
+# from .linear_attention.linear_attn import LinearAttentionFusedChunkKernel, LinearAttentionFusedRecurrentKernel
+from .bitnet import Bitnet_158_int8xint2_kernel
+
+__all__ = [
+    "MHAKernel", "MHADecodeKernel", "MLAKernel", "GQAKernel", "GQADecodeKernel",
+    "MambaChunkScanKernel", "MambaChunkStateKernel", "BlockSparseAttentionKernel",
+    # "LinearAttentionFusedChunkKernel", "LinearAttentionFusedRecurrentKernel",
+    "Bitnet_158_int8xint2_kernel"
+]