tile-kernels 1.0.0__py3-none-any.whl

tile_kernels/__init__.py

@@ -0,0 +1,15 @@
+import tilelang
+
+from . import (
+    config,
+    engram,
+    mhc,
+    modeling,
+    moe,
+    quant,
+    transpose,
+    torch,
+    testing,
+)
+
+from .config import get_num_sms, get_device_num_sms, set_num_sms

tile_kernels/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '1.0.0'
+__version_tuple__ = version_tuple = (1, 0, 0)
+
+__commit_id__ = commit_id = None

tile_kernels/config.py

@@ -0,0 +1,29 @@
+import functools
+import torch
+
+_num_sms = 0
+
+
+@functools.lru_cache(maxsize=None)
+def get_device_num_sms() -> int:
+    prop = torch.cuda.get_device_properties(torch.cuda.current_device())
+    return prop.multi_processor_count
+
+
+def set_num_sms(num_sms: int) -> None:
+    global _num_sms
+    assert 0 < num_sms <= get_device_num_sms()
+    _num_sms = num_sms
+
+
+def get_num_sms() -> int:
+    global _num_sms
+    if _num_sms == 0:
+        return get_device_num_sms()
+    return _num_sms
+
+
+@functools.lru_cache(maxsize=None)
+def get_max_smem_per_sm() -> int:
+    prop = torch.cuda.get_device_properties(torch.cuda.current_device())
+    return prop.shared_memory_per_multiprocessor

tile_kernels/engram/__init__.py

@@ -0,0 +1,4 @@
+from .engram_fused_weight_kernel import fused_weight
+from .engram_gate_kernel import engram_gate_fwd, engram_gate_bwd
+from .engram_grad_w_reduce_kernel import grad_w_reduce
+from .engram_hash_kernel import engram_hash

tile_kernels/engram/engram_fused_weight_kernel.py

@@ -0,0 +1,59 @@
+import os
+
+import torch
+import tilelang
+from tilelang import language as T
+
+
+@tilelang.jit(
+    pass_configs={
+        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+    },
+)
+def get_engram_fused_weight_kernel(hidden_size: int, hc_mult: int):
+    """Elementwise bf16 x bf16 -> fp32 for weight_hidden * weight_embed."""
+    threads = 32
+    vec_size = 8
+    blk_d = threads * vec_size
+    assert hidden_size % blk_d == 0
+    num_blk = hidden_size // blk_d
+
+    @T.prim_func
+    def engram_fused_weight_kernel(
+        weight_hidden: T.Tensor[(hc_mult, hidden_size), T.bfloat16],
+        weight_embed: T.Tensor[(hc_mult, hidden_size), T.bfloat16],
+        weight_fused: T.Tensor[(hc_mult, hidden_size), T.float],
+    ):
+        with T.Kernel(hc_mult, num_blk, threads=threads) as (pid_h, pid_b):
+            tid = T.get_thread_binding()
+            a_local = T.alloc_local((vec_size,), T.float)
+            b_local = T.alloc_local((vec_size,), T.float)
+            for i_k in T.vectorized(vec_size):
+                a_local[i_k] = weight_hidden[pid_h, pid_b * blk_d + tid * vec_size + i_k]
+                b_local[i_k] = weight_embed[pid_h, pid_b * blk_d + tid * vec_size + i_k]
+            for i_k in T.vectorized(vec_size):
+                weight_fused[pid_h, pid_b * blk_d + tid * vec_size + i_k] = a_local[i_k] * b_local[i_k]
+
+    return engram_fused_weight_kernel
+
+
+def fused_weight(weight_hidden: torch.Tensor, weight_embed: torch.Tensor) -> torch.Tensor:
+    """Compute weight_hidden * weight_embed in fp32.
+
+    Args:
+        weight_hidden: Shape (hc_mult, hidden_size), bfloat16.
+        weight_embed: Shape (hc_mult, hidden_size), bfloat16.
+
+    Returns:
+        weight_fused: Shape (hc_mult, hidden_size), float32.
+    """
+    hc_mult, hidden_size = weight_hidden.shape
+
+    kernel = get_engram_fused_weight_kernel(hidden_size, hc_mult)
+    if int(os.getenv('TK_PRINT_KERNEL_SOURCE', 0)):
+        print(kernel.get_kernel_source())
+
+    weight_fused = torch.empty(hc_mult, hidden_size, dtype=torch.float32, device=weight_hidden.device)
+    kernel(weight_hidden, weight_embed, weight_fused)
+
+    return weight_fused