Uploaded using `kernel-builder`.

Browse files

Files changed (13) hide show

build/torch-cuda/__init__.py +68 -0
build/torch-cuda/_errors.py +23 -0
build/torch-cuda/_ops.py +38 -0
build/torch-cuda/_runtime.py +208 -0
build/torch-cuda/_version.py +22 -0
build/torch-cuda/kernels/__init__.py +7 -0
build/torch-cuda/kernels/compile_cache.py +37 -0
build/torch-cuda/kernels/flash_fwd.py +978 -0
build/torch-cuda/metadata.json +10 -0
build/torch-cuda/sdpa.py +54 -0
build/torch-cuda/tiledattention/__init__.py +26 -0
build/torch-cuda/utils/__init__.py +7 -0
build/torch-cuda/utils/checks.py +58 -0

build/torch-cuda/__init__.py ADDED Viewed

	@@ -0,0 +1,68 @@

+"""Public package interface for tiledattention."""
+from __future__ import annotations
+from typing import TYPE_CHECKING
+try:
+    from ._errors import (
+        DependencyError,
+        DTypeNotSupportedError,
+        InvalidShapeError,
+        TiledAttentionError,
+        UnsupportedPlatformError,
+    )
+    from ._version import __version__
+except ImportError:  # pragma: no cover - flattened kernel package layout
+    from _errors import (  # type: ignore
+        DependencyError,
+        DTypeNotSupportedError,
+        InvalidShapeError,
+        TiledAttentionError,
+        UnsupportedPlatformError,
+    )
+    from _version import __version__  # type: ignore
+if TYPE_CHECKING:
+    import torch
+def sdpa(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    causal: bool = False,
+    scale: float | None = None,
+) -> torch.Tensor:
+    """
+    Compute scaled dot-product attention.
+    It is part of the public SDPA execution path.
+    Args:
+        q: Query tensor in attention layout.
+        k: Key tensor in attention layout.
+        v: Value tensor in attention layout.
+        causal: Whether causal masking is enabled.
+        scale: Attention scaling factor.
+    Returns:
+        torch.Tensor: Function result value.
+    """
+    try:
+        from .sdpa import sdpa as _sdpa
+    except ImportError:  # pragma: no cover - flattened kernel package layout
+        from sdpa import sdpa as _sdpa  # type: ignore
+    return _sdpa(q, k, v, causal=causal, scale=scale)
+__all__ = [
+    "__version__",
+    "sdpa",
+    "TiledAttentionError",
+    "UnsupportedPlatformError",
+    "InvalidShapeError",
+    "DTypeNotSupportedError",
+    "DependencyError",
+]

build/torch-cuda/_errors.py ADDED Viewed

	@@ -0,0 +1,23 @@

+"""Custom exception types for tiledattention."""
+from __future__ import annotations
+class TiledAttentionError(Exception):
+    """Base exception for tiledattention."""
+class UnsupportedPlatformError(TiledAttentionError):
+    """Raised when runtime platform requirements are not satisfied."""
+class InvalidShapeError(TiledAttentionError):
+    """Raised when q/k/v tensors do not match expected SDPA input contracts."""
+class DTypeNotSupportedError(TiledAttentionError):
+    """Raised when input tensor dtypes are unsupported."""
+class DependencyError(TiledAttentionError):
+    """Raised when a required dependency is missing."""

build/torch-cuda/_ops.py ADDED Viewed

	@@ -0,0 +1,38 @@

+import torch
+def get_backend() -> str:
+    """Detect the backend by inspecting torch."""
+    import torch
+    if hasattr(torch, "neuron"):
+        # Needs to be sorted before specific Torch builds, since Neuron
+        # extension can be loaded into e.g. CUDA Torch builds.
+        return "neuron"
+    elif torch.version.cuda is not None:
+        return "cuda"
+    elif torch.version.hip is not None:
+        return "rocm"
+    elif torch.backends.mps.is_available():
+        return "metal"
+    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+        return "xpu"
+    else:
+        return "cpu"
+def _find_ops_name() -> str:
+    kernel_name = "tiledattention"
+    unique_id = "703d09c_dirty"
+    backend = get_backend()
+    return f"_{kernel_name}_{backend}_{unique_id}"
+_OPS_NAME = _find_ops_name()
+ops = getattr(torch.ops, _OPS_NAME)
+def add_op_namespace_prefix(op_name: str) -> str:
+    """
+    Prefix op by namespace.
+    """
+    return f"{_OPS_NAME}::{op_name}"

build/torch-cuda/_runtime.py ADDED Viewed

	@@ -0,0 +1,208 @@

+"""Runtime checks and environment introspection for tiledattention."""
+from __future__ import annotations
+import importlib
+import os
+from threading import Lock
+from types import ModuleType
+try:
+    from ._errors import DependencyError, UnsupportedPlatformError
+except ImportError:  # pragma: no cover - flattened kernel package layout
+    from _errors import DependencyError, UnsupportedPlatformError  # type: ignore
+_SUPPORTED_CC_MAJORS = {10, 12}
+_runtime_ready = False
+_runtime_lock = Lock()
+_cached_torch: ModuleType | None = None
+_cached_cupy: ModuleType | None = None
+_cached_cutile: ModuleType | None = None
+def _get_torch_module() -> ModuleType:
+    """
+    Get torch module.
+    It supports runtime dependency loading and platform validation.
+    Returns:
+        ModuleType: Function result value.
+    """
+    global _cached_torch
+    if _cached_torch is not None:
+        return _cached_torch
+    try:
+        _cached_torch = importlib.import_module("torch")
+        return _cached_torch
+    except ModuleNotFoundError as exc:
+        raise DependencyError(
+            "PyTorch is required for tiledattention. Install a CUDA-enabled torch build first."
+        ) from exc
+def get_torch_module() -> ModuleType:
+    """
+    Get torch module.
+    It supports runtime dependency loading and platform validation.
+    Returns:
+        ModuleType: Function result value.
+    """
+    return _get_torch_module()
+def _get_cupy_module() -> ModuleType:
+    """
+    Get cupy module.
+    It supports runtime dependency loading and platform validation.
+    Returns:
+        ModuleType: Function result value.
+    """
+    global _cached_cupy
+    if _cached_cupy is not None:
+        return _cached_cupy
+    try:
+        _cached_cupy = importlib.import_module("cupy")
+        return _cached_cupy
+    except ModuleNotFoundError as exc:
+        raise DependencyError(
+            "CuPy is required at runtime. For CUDA 13 environments, install a compatible wheel (e.g., pip install cupy-cuda13x)."
+        ) from exc
+def get_cupy_module() -> ModuleType:
+    """
+    Get cupy module.
+    It supports runtime dependency loading and platform validation.
+    Returns:
+        ModuleType: Function result value.
+    """
+    return _get_cupy_module()
+def _cutile_candidates() -> tuple[str, ...]:
+    """
+    Internal helper for cutile candidates.
+    It supports runtime dependency loading and platform validation.
+    Returns:
+        tuple[str, ...]: Function result value.
+    """
+    raw = os.getenv("TILEDATTENTION_CUTILE_MODULE", "").strip()
+    if raw:
+        parsed = tuple(part.strip() for part in raw.split(",") if part.strip())
+        if parsed:
+            return parsed
+    return ("cutile", "cuda.tile", "cuda_tile")
+def _get_cutile_module() -> ModuleType:
+    """
+    Get cutile module.
+    It supports runtime dependency loading and platform validation.
+    Returns:
+        ModuleType: Function result value.
+    """
+    global _cached_cutile
+    if _cached_cutile is not None:
+        return _cached_cutile
+    candidates = _cutile_candidates()
+    for module_name in candidates:
+        try:
+            _cached_cutile = importlib.import_module(module_name)
+            return _cached_cutile
+        except ModuleNotFoundError:
+            continue
+    candidate_text = ", ".join(candidates)
+    raise DependencyError(
+        "cuTile Python module could not be imported. "
+        f"Tried: {candidate_text}. "
+        "Install NVIDIA CUDA Tile Python (e.g., pip install cuda-tile) or set "
+        "TILEDATTENTION_CUTILE_MODULE=<module_name>."
+    )
+def get_cutile_module() -> ModuleType:
+    """
+    Get cutile module.
+    It supports runtime dependency loading and platform validation.
+    Returns:
+        ModuleType: Function result value.
+    """
+    return _get_cutile_module()
+def _query_compute_capability(torch_mod: ModuleType) -> tuple[int, int]:
+    """
+    Internal helper for query compute capability.
+    It supports runtime dependency loading and platform validation.
+    Args:
+        torch_mod: Imported torch module instance.
+    Returns:
+        tuple[int, int]: Function result value.
+    """
+    try:
+        device_index = torch_mod.cuda.current_device()
+        major, minor = torch_mod.cuda.get_device_capability(device_index)
+        return int(major), int(minor)
+    except Exception as exc:  # pragma: no cover - defensive path
+        raise UnsupportedPlatformError(
+            "Unable to query GPU compute capability from torch.cuda."
+        ) from exc
+def require_supported_runtime() -> None:
+    """
+    Run require supported runtime.
+    It supports runtime dependency loading and platform validation.
+    """
+    global _runtime_ready
+    with _runtime_lock:
+        if _runtime_ready:
+            return
+        torch_mod = _get_torch_module()
+        if not bool(torch_mod.cuda.is_available()):
+            raise UnsupportedPlatformError(
+                "tiledattention requires CUDA + Blackwell GPU. "
+                "torch.cuda.is_available() is False."
+            )
+        major, minor = _query_compute_capability(torch_mod)
+        if major not in _SUPPORTED_CC_MAJORS:
+            raise UnsupportedPlatformError(
+                f"Unsupported GPU: compute capability {major}.{minor}; "
+                "requires Blackwell-class GPU (10.x or 12.x)."
+            )
+        _get_cupy_module()
+        _get_cutile_module()
+        _runtime_ready = True
+def _reset_runtime_cache_for_tests() -> None:
+    """
+    Internal helper for reset runtime cache for tests.
+    It supports runtime dependency loading and platform validation.
+    """
+    global _runtime_ready, _cached_torch, _cached_cupy, _cached_cutile
+    with _runtime_lock:
+        _runtime_ready = False
+        _cached_torch = None
+        _cached_cupy = None
+        _cached_cutile = None

build/torch-cuda/_version.py ADDED Viewed

	@@ -0,0 +1,22 @@

+"""Version helpers for tiledattention."""
+from __future__ import annotations
+from importlib.metadata import PackageNotFoundError, version
+def _resolve_version() -> str:
+    """
+    Resolve version.
+    It is used by the tiledattention runtime and tooling.
+    Returns:
+        str: Function result value.
+    """
+    try:
+        return version("tiledattention")
+    except PackageNotFoundError:
+        return "0.0.0"
+__version__ = _resolve_version()

build/torch-cuda/kernels/__init__.py ADDED Viewed

	@@ -0,0 +1,7 @@

+"""Kernel entrypoints."""
+from __future__ import annotations
+from .flash_fwd import run_flash_fwd
+__all__ = ["run_flash_fwd"]

build/torch-cuda/kernels/compile_cache.py ADDED Viewed

	@@ -0,0 +1,37 @@

+"""In-memory kernel cache for cuTile kernel factories."""
+from __future__ import annotations
+from collections.abc import Callable
+KernelCallable = Callable[..., object]
+KernelKey = tuple[object, ...]
+_KERNEL_CACHE: dict[KernelKey, KernelCallable] = {}
+def get_kernel(key: KernelKey, factory: Callable[[], KernelCallable]) -> KernelCallable:
+    """
+    Get kernel.
+    It manages in-memory reuse of compiled kernel callables.
+    Args:
+        key: Cache key for kernel lookup.
+        factory: Factory that creates a kernel callable.
+    Returns:
+        KernelCallable: Function result value.
+    """
+    kernel = _KERNEL_CACHE.get(key)
+    if kernel is None:
+        kernel = factory()
+        _KERNEL_CACHE[key] = kernel
+    return kernel
+def reset_cache_for_tests() -> None:
+    """
+    Run reset cache for tests.
+    It manages in-memory reuse of compiled kernel callables.
+    """
+    _KERNEL_CACHE.clear()

build/torch-cuda/kernels/flash_fwd.py ADDED Viewed

	@@ -0,0 +1,978 @@

+"""FlashAttention-style forward kernel launcher."""
+from __future__ import annotations
+import os
+from typing import Any
+try:
+    from .. import _runtime
+    from .._errors import DTypeNotSupportedError
+except ImportError:  # pragma: no cover - flattened kernel package layout
+    import _runtime  # type: ignore
+    from _errors import DTypeNotSupportedError  # type: ignore
+from .compile_cache import get_kernel
+_DEFAULT_TM = 64
+_DEFAULT_TN = 64
+_NEG_LARGE = -1.0e30
+_DIRECT_HEAD_DIMS = frozenset({64, 128})
+_ALIGNED_FASTPATH_HEAD_DIMS = frozenset({64, 128})
+_ALIGNED_FASTPATH_MIN_SEQ_BY_HEAD_DIM = {
+    64: 1024,
+    128: 2048,
+}
+_CHUNKED_HEAD_DIM_PARTS = {
+    96: (64, 32),
+    160: (128, 32),
+}
+def _ct_dtype_for_torch_dtype(ct: Any, torch_dtype: Any) -> Any:
+    """
+    Internal helper for ct dtype for torch dtype.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        ct: Imported cuTile module instance.
+        torch_dtype: Torch dtype value to map into cuTile dtype.
+    Returns:
+        Any: Function result value.
+    """
+    torch_mod = _runtime.get_torch_module()
+    if torch_dtype == torch_mod.float16:
+        return ct.float16
+    if torch_dtype == torch_mod.bfloat16:
+        return ct.bfloat16
+    raise DTypeNotSupportedError("q dtype must be torch.float16 or torch.bfloat16.")
+def _get_env_int(name: str, default: int | None) -> int | None:
+    """
+    Get env int.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        name: Identifier or metric name.
+        default: Default value used when environment is unset.
+    Returns:
+        int | None: Function result value.
+    """
+    raw = os.getenv(name)
+    if raw is None or raw.strip() == "":
+        return default
+    try:
+        return int(raw)
+    except ValueError as exc:  # pragma: no cover - defensive
+        raise ValueError(f"Environment variable {name} must be an integer.") from exc
+def _resolve_tile_config() -> tuple[int, int]:
+    """
+    Resolve tile config.
+    It is used during kernel configuration, compilation, or launch.
+    Returns:
+        tuple[int, int]: Function result value.
+    """
+    tile_m = _get_env_int("TILEDATTN_TILE_M", None)
+    tile_n = _get_env_int("TILEDATTN_TILE_N", None)
+    if tile_m is None:
+        tile_m = _DEFAULT_TM
+    if tile_n is None:
+        tile_n = _DEFAULT_TN
+    if tile_m <= 0 or tile_n <= 0:
+        raise ValueError("TILEDATTN_TILE_M and TILEDATTN_TILE_N must be positive.")
+    return tile_m, tile_n
+def _resolve_kernel_options() -> tuple[int, int | None, int | None]:
+    """
+    Resolve kernel options.
+    It is used during kernel configuration, compilation, or launch.
+    Returns:
+        tuple[int, int | None, int | None]: Function result value.
+    """
+    opt_level = _get_env_int("TILEDATTN_KERNEL_OPT_LEVEL", 3)
+    occupancy = _get_env_int("TILEDATTN_KERNEL_OCCUPANCY", None)
+    num_ctas = _get_env_int("TILEDATTN_KERNEL_NUM_CTAS", None)
+    assert opt_level is not None
+    return opt_level, occupancy, num_ctas
+def _default_accum_mode_for_shape(*, seq_len: int, head_dim: int, causal: bool) -> str:
+    # Empirical policy from reduced benchmark:
+    # - D=64 long-ish non-causal shapes prefer fp16 accumulation.
+    # - D=128 long non-causal shapes prefer fp32 accumulation.
+    """
+    Internal helper for default accum mode for shape.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        seq_len: Sequence length used for heuristic decisions.
+        head_dim: Attention head dimension.
+        causal: Whether causal masking is enabled.
+    Returns:
+        str: Function result value.
+    """
+    if (not causal) and head_dim == 64 and seq_len >= 1024:
+        return "fp16"
+    return "fp32"
+def _resolve_accum_mode(*, seq_len: int, head_dim: int, causal: bool) -> str:
+    """
+    Resolve accum mode.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        seq_len: Sequence length used for heuristic decisions.
+        head_dim: Attention head dimension.
+        causal: Whether causal masking is enabled.
+    Returns:
+        str: Function result value.
+    """
+    raw_mode = os.getenv("TILEDATTN_ACCUM_MODE")
+    if raw_mode is None or raw_mode.strip() == "":
+        return _default_accum_mode_for_shape(seq_len=seq_len, head_dim=head_dim, causal=causal)
+    mode = raw_mode.strip().lower()
+    if mode == "auto":
+        return _default_accum_mode_for_shape(seq_len=seq_len, head_dim=head_dim, causal=causal)
+    if mode not in {"fp32", "fp16"}:
+        raise ValueError(
+            f"Unsupported TILEDATTN_ACCUM_MODE={mode!r}. Use one of: auto, fp32, fp16."
+        )
+    return mode
+def _resolve_kernel_head_dim(head_dim: int) -> tuple[int, int]:
+    """
+    Resolve kernel head dim.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        head_dim: Attention head dimension.
+    Returns:
+        tuple[int, int]: Function result value.
+    """
+    if head_dim in _DIRECT_HEAD_DIMS:
+        return head_dim, 0
+    kernel_head_dim = 1 << (head_dim - 1).bit_length()
+    return kernel_head_dim, kernel_head_dim - head_dim
+def _resolve_chunk_plan(head_dim: int) -> tuple[tuple[int, int], ...] | None:
+    """
+    Resolve chunk plan.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        head_dim: Attention head dimension.
+    Returns:
+        tuple[tuple[int, int], ...] | None: Function result value.
+    """
+    enabled = os.getenv("TILEDATTN_CHUNKED_HEAD_DIMS", "").strip()
+    if enabled == "":
+        return None
+    try:
+        enabled_dims = {int(x.strip()) for x in enabled.split(",") if x.strip() != ""}
+    except ValueError as exc:
+        raise ValueError(
+            "Environment variable TILEDATTN_CHUNKED_HEAD_DIMS must be a comma-separated list of integers."
+        ) from exc
+    if head_dim not in enabled_dims:
+        return None
+    parts = _CHUNKED_HEAD_DIM_PARTS.get(head_dim)
+    if parts is None:
+        return None
+    offset = 0
+    plan: list[tuple[int, int]] = []
+    for width in parts:
+        plan.append((offset, width))
+        offset += width
+    return tuple(plan)
+def _should_use_aligned_noncausal_fastpath(
+    *,
+    seq_len: int,
+    head_dim: int,
+    causal: bool,
+    pad_dim: int,
+    chunk_plan: tuple[tuple[int, int], ...] | None,
+    tile_m: int,
+    tile_n: int,
+) -> bool:
+    """
+    Internal helper for should use aligned noncausal fastpath.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        seq_len: Sequence length used for heuristic decisions.
+        head_dim: Attention head dimension.
+        causal: Whether causal masking is enabled.
+        pad_dim: Function argument.
+        chunk_plan: Head-dimension chunk decomposition plan.
+        tile_m: Tile size in the query-row dimension.
+        tile_n: Tile size in the key/value-column dimension.
+    Returns:
+        bool: Function result value.
+    """
+    if os.getenv("TILEDATTN_DISABLE_ALIGNED_FASTPATH", "").strip() not in {"", "0", "false", "False"}:
+        return False
+    if causal:
+        return False
+    if chunk_plan is not None:
+        return False
+    if pad_dim != 0:
+        return False
+    if head_dim not in _ALIGNED_FASTPATH_HEAD_DIMS:
+        return False
+    min_seq = _ALIGNED_FASTPATH_MIN_SEQ_BY_HEAD_DIM.get(head_dim, 2048)
+    if seq_len < min_seq:
+        return False
+    return (seq_len % tile_m == 0) and (seq_len % tile_n == 0)
+def _default_tile_config_for_shape(*, seq_len: int, head_dim: int, causal: bool) -> tuple[int, int]:
+    """
+    Internal helper for default tile config for shape.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        seq_len: Sequence length used for heuristic decisions.
+        head_dim: Attention head dimension.
+        causal: Whether causal masking is enabled.
+    Returns:
+        tuple[int, int]: Function result value.
+    """
+    tile_m = _DEFAULT_TM
+    tile_n = _DEFAULT_TN
+    # Mid/long non-causal D=128 favors wider N tiles in current tuning results.
+    if (not causal) and head_dim == 128 and seq_len >= 2048:
+        tile_n = 128
+    return tile_m, tile_n
+def make_flashattn_fwd_kernel(
+    tile_m: int,
+    tile_n: int,
+    head_dim: int,
+    *,
+    dtype: Any,
+    causal: bool,
+    accum_mode: str,
+    opt_level: int,
+    occupancy: int | None,
+    num_ctas: int | None,
+):
+    """
+    Create flashattn fwd kernel.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        tile_m: Tile size in the query-row dimension.
+        tile_n: Tile size in the key/value-column dimension.
+        head_dim: Attention head dimension.
+        dtype: Target data type used for computation.
+        causal: Whether causal masking is enabled.
+        accum_mode: Function argument.
+        opt_level: Kernel optimization level forwarded to cuTile.
+        occupancy: Optional occupancy hint for kernel launch.
+        num_ctas: Optional CTA count hint for kernel launch.
+    Returns:
+        object: Function result value.
+    """
+    ct = _runtime.get_cutile_module()
+    kernel_kwargs: dict[str, Any] = {"opt_level": opt_level}
+    if occupancy is not None:
+        kernel_kwargs["occupancy"] = occupancy
+    if num_ctas is not None:
+        kernel_kwargs["num_ctas"] = num_ctas
+    if accum_mode == "fp16":
+        @ct.kernel(**kernel_kwargs)
+        def flash_fwd_kernel_fp16acc(q, k_t, v, out, scale):
+            """
+            Run flash fwd kernel fp16acc.
+            It is used during kernel configuration, compilation, or launch.
+            Args:
+                q: Query tensor in attention layout.
+                k_t: Transposed key tensor view used by the kernel.
+                v: Value tensor in attention layout.
+                out: Output tensor buffer.
+                scale: Attention scaling factor.
+            """
+            bh_idx = ct.bid(0)
+            q_tile_idx = ct.bid(1)
+            # Keep GEMM inputs in low precision to allow tensor-core codegen.
+            q_tile = ct.load(q, index=(bh_idx, q_tile_idx, 0), shape=(1, tile_m, head_dim))
+            seq_len = q.shape[1]
+            num_k_tiles = ct.cdiv(seq_len, tile_n)
+            if causal:
+                # For causal mode, this query tile never attends beyond its last row index.
+                causal_cols = ct.minimum((q_tile_idx + 1) * tile_m, seq_len)
+                num_k_tiles = ct.cdiv(causal_cols, tile_n)
+            row = q_tile_idx * tile_m + ct.arange(tile_m, dtype=ct.int32)
+            row = ct.expand_dims(row, 1)
+            row = ct.expand_dims(row, 0)
+            row_in_bounds = row < seq_len
+            m_i = ct.full((1, tile_m, 1), _NEG_LARGE, ct.float32)
+            l_i = ct.zeros((1, tile_m, 1), ct.float32)
+            acc = ct.zeros((1, tile_m, head_dim), dtype)
+            for k_tile_idx in range(num_k_tiles):
+                k_tile_t = ct.load(k_t, index=(bh_idx, 0, k_tile_idx), shape=(1, head_dim, tile_n))
+                score = ct.matmul(q_tile, k_tile_t)
+                score = ct.astype(score, ct.float32) * scale
+                col = k_tile_idx * tile_n + ct.arange(tile_n, dtype=ct.int32)
+                col = ct.expand_dims(col, 0)
+                col = ct.expand_dims(col, 0)
+                key_in_bounds = col < seq_len
+                valid = row_in_bounds & key_in_bounds
+                if causal:
+                    valid = valid & (row >= col)
+                score = ct.where(valid, score, _NEG_LARGE)
+                tile_max = ct.max(score, axis=2, keepdims=True)
+                m_next = ct.maximum(m_i, tile_max)
+                alpha = ct.exp(m_i - m_next)
+                p = ct.exp(score - m_next)
+                l_i = l_i * alpha + ct.sum(p, axis=2, keepdims=True)
+                v_tile = ct.load(v, index=(bh_idx, k_tile_idx, 0), shape=(1, tile_n, head_dim))
+                p_lowp = ct.astype(p, dtype)
+                alpha_lowp = ct.astype(alpha, dtype)
+                acc = acc * alpha_lowp + ct.matmul(p_lowp, v_tile)
+                m_i = m_next
+            safe_l = ct.where(row_in_bounds, l_i, 1.0)
+            # Row-wise reciprocal avoids a full per-element divide over head_dim.
+            inv_l = 1.0 / safe_l
+            out_tile = acc * ct.astype(inv_l, dtype)
+            out_tile = ct.where(row_in_bounds, out_tile, 0.0)
+            ct.store(out, index=(bh_idx, q_tile_idx, 0), tile=out_tile)
+        return flash_fwd_kernel_fp16acc
+    @ct.kernel(**kernel_kwargs)
+    def flash_fwd_kernel_fp32acc(q, k_t, v, out, scale):
+        """
+        Run flash fwd kernel fp32acc.
+        It is used during kernel configuration, compilation, or launch.
+        Args:
+            q: Query tensor in attention layout.
+            k_t: Transposed key tensor view used by the kernel.
+            v: Value tensor in attention layout.
+            out: Output tensor buffer.
+            scale: Attention scaling factor.
+        """
+        bh_idx = ct.bid(0)
+        q_tile_idx = ct.bid(1)
+        # Keep GEMM inputs in low precision to allow tensor-core codegen.
+        q_tile = ct.load(q, index=(bh_idx, q_tile_idx, 0), shape=(1, tile_m, head_dim))
+        seq_len = q.shape[1]
+        num_k_tiles = ct.cdiv(seq_len, tile_n)
+        if causal:
+            # For causal mode, this query tile never attends beyond its last row index.
+            causal_cols = ct.minimum((q_tile_idx + 1) * tile_m, seq_len)
+            num_k_tiles = ct.cdiv(causal_cols, tile_n)
+        row = q_tile_idx * tile_m + ct.arange(tile_m, dtype=ct.int32)
+        row = ct.expand_dims(row, 1)
+        row = ct.expand_dims(row, 0)
+        row_in_bounds = row < seq_len
+        m_i = ct.full((1, tile_m, 1), _NEG_LARGE, ct.float32)
+        l_i = ct.zeros((1, tile_m, 1), ct.float32)
+        acc = ct.zeros((1, tile_m, head_dim), ct.float32)
+        for k_tile_idx in range(num_k_tiles):
+            k_tile_t = ct.load(k_t, index=(bh_idx, 0, k_tile_idx), shape=(1, head_dim, tile_n))
+            score = ct.matmul(q_tile, k_tile_t)
+            score = ct.astype(score, ct.float32) * scale
+            col = k_tile_idx * tile_n + ct.arange(tile_n, dtype=ct.int32)
+            col = ct.expand_dims(col, 0)
+            col = ct.expand_dims(col, 0)
+            key_in_bounds = col < seq_len
+            valid = row_in_bounds & key_in_bounds
+            if causal:
+                valid = valid & (row >= col)
+            score = ct.where(valid, score, _NEG_LARGE)
+            tile_max = ct.max(score, axis=2, keepdims=True)
+            m_next = ct.maximum(m_i, tile_max)
+            alpha = ct.exp(m_i - m_next)
+            p = ct.exp(score - m_next)
+            l_i = l_i * alpha + ct.sum(p, axis=2, keepdims=True)
+            v_tile = ct.load(v, index=(bh_idx, k_tile_idx, 0), shape=(1, tile_n, head_dim))
+            p_lowp = ct.astype(p, dtype)
+            acc_update = ct.matmul(p_lowp, v_tile)
+            acc = acc * alpha + ct.astype(acc_update, ct.float32)
+            m_i = m_next
+        safe_l = ct.where(row_in_bounds, l_i, 1.0)
+        # Row-wise reciprocal avoids a full per-element divide over head_dim.
+        inv_l = 1.0 / safe_l
+        out_tile = acc * inv_l
+        out_tile = ct.where(row_in_bounds, out_tile, 0.0)
+        out_tile = ct.astype(out_tile, dtype)
+        ct.store(out, index=(bh_idx, q_tile_idx, 0), tile=out_tile)
+    return flash_fwd_kernel_fp32acc
+def make_flashattn_fwd_kernel_aligned_noncausal(
+    tile_m: int,
+    tile_n: int,
+    head_dim: int,
+    *,
+    dtype: Any,
+    accum_mode: str,
+    opt_level: int,
+    occupancy: int | None,
+    num_ctas: int | None,
+):
+    """
+    Create flashattn fwd kernel aligned noncausal.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        tile_m: Tile size in the query-row dimension.
+        tile_n: Tile size in the key/value-column dimension.
+        head_dim: Attention head dimension.
+        dtype: Target data type used for computation.
+        accum_mode: Function argument.
+        opt_level: Kernel optimization level forwarded to cuTile.
+        occupancy: Optional occupancy hint for kernel launch.
+        num_ctas: Optional CTA count hint for kernel launch.
+    Returns:
+        object: Function result value.
+    """
+    ct = _runtime.get_cutile_module()
+    kernel_kwargs: dict[str, Any] = {"opt_level": opt_level}
+    if occupancy is not None:
+        kernel_kwargs["occupancy"] = occupancy
+    if num_ctas is not None:
+        kernel_kwargs["num_ctas"] = num_ctas
+    if accum_mode == "fp16":
+        @ct.kernel(**kernel_kwargs)
+        def flash_fwd_kernel_fp16acc_aligned(q, k_t, v, out, scale):
+            """
+            Run flash fwd kernel fp16acc aligned.
+            It is used during kernel configuration, compilation, or launch.
+            Args:
+                q: Query tensor in attention layout.
+                k_t: Transposed key tensor view used by the kernel.
+                v: Value tensor in attention layout.
+                out: Output tensor buffer.
+                scale: Attention scaling factor.
+            """
+            bh_idx = ct.bid(0)
+            q_tile_idx = ct.bid(1)
+            q_tile = ct.load(q, index=(bh_idx, q_tile_idx, 0), shape=(1, tile_m, head_dim))
+            seq_len = q.shape[1]
+            num_k_tiles = seq_len // tile_n
+            m_i = ct.full((1, tile_m, 1), _NEG_LARGE, ct.float32)
+            l_i = ct.zeros((1, tile_m, 1), ct.float32)
+            acc = ct.zeros((1, tile_m, head_dim), dtype)
+            for k_tile_idx in range(num_k_tiles):
+                k_tile_t = ct.load(k_t, index=(bh_idx, 0, k_tile_idx), shape=(1, head_dim, tile_n))
+                score = ct.astype(ct.matmul(q_tile, k_tile_t), ct.float32) * scale
+                tile_max = ct.max(score, axis=2, keepdims=True)
+                m_next = ct.maximum(m_i, tile_max)
+                alpha = ct.exp(m_i - m_next)
+                p = ct.exp(score - m_next)
+                l_i = l_i * alpha + ct.sum(p, axis=2, keepdims=True)
+                v_tile = ct.load(v, index=(bh_idx, k_tile_idx, 0), shape=(1, tile_n, head_dim))
+                p_lowp = ct.astype(p, dtype)
+                alpha_lowp = ct.astype(alpha, dtype)
+                acc = acc * alpha_lowp + ct.matmul(p_lowp, v_tile)
+                m_i = m_next
+            inv_l = 1.0 / l_i
+            out_tile = acc * ct.astype(inv_l, dtype)
+            ct.store(out, index=(bh_idx, q_tile_idx, 0), tile=out_tile)
+        return flash_fwd_kernel_fp16acc_aligned
+    @ct.kernel(**kernel_kwargs)
+    def flash_fwd_kernel_fp32acc_aligned(q, k_t, v, out, scale):
+        """
+        Run flash fwd kernel fp32acc aligned.
+        It is used during kernel configuration, compilation, or launch.
+        Args:
+            q: Query tensor in attention layout.
+            k_t: Transposed key tensor view used by the kernel.
+            v: Value tensor in attention layout.
+            out: Output tensor buffer.
+            scale: Attention scaling factor.
+        """
+        bh_idx = ct.bid(0)
+        q_tile_idx = ct.bid(1)
+        q_tile = ct.load(q, index=(bh_idx, q_tile_idx, 0), shape=(1, tile_m, head_dim))
+        seq_len = q.shape[1]
+        num_k_tiles = seq_len // tile_n
+        m_i = ct.full((1, tile_m, 1), _NEG_LARGE, ct.float32)
+        l_i = ct.zeros((1, tile_m, 1), ct.float32)
+        acc = ct.zeros((1, tile_m, head_dim), ct.float32)
+        for k_tile_idx in range(num_k_tiles):
+            k_tile_t = ct.load(k_t, index=(bh_idx, 0, k_tile_idx), shape=(1, head_dim, tile_n))
+            score = ct.astype(ct.matmul(q_tile, k_tile_t), ct.float32) * scale
+            tile_max = ct.max(score, axis=2, keepdims=True)
+            m_next = ct.maximum(m_i, tile_max)
+            alpha = ct.exp(m_i - m_next)
+            p = ct.exp(score - m_next)
+            l_i = l_i * alpha + ct.sum(p, axis=2, keepdims=True)
+            v_tile = ct.load(v, index=(bh_idx, k_tile_idx, 0), shape=(1, tile_n, head_dim))
+            p_lowp = ct.astype(p, dtype)
+            acc = acc * alpha + ct.astype(ct.matmul(p_lowp, v_tile), ct.float32)
+            m_i = m_next
+        out_tile = ct.astype(acc * (1.0 / l_i), dtype)
+        ct.store(out, index=(bh_idx, q_tile_idx, 0), tile=out_tile)
+    return flash_fwd_kernel_fp32acc_aligned
+def make_flashattn_fwd_kernel_chunked(
+    tile_m: int,
+    tile_n: int,
+    head_dim: int,
+    *,
+    chunk_plan: tuple[tuple[int, int], ...],
+    dtype: Any,
+    causal: bool,
+    accum_mode: str,
+    opt_level: int,
+    occupancy: int | None,
+    num_ctas: int | None,
+):
+    """
+    Create flashattn fwd kernel chunked.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        tile_m: Tile size in the query-row dimension.
+        tile_n: Tile size in the key/value-column dimension.
+        head_dim: Attention head dimension.
+        chunk_plan: Head-dimension chunk decomposition plan.
+        dtype: Target data type used for computation.
+        causal: Whether causal masking is enabled.
+        accum_mode: Function argument.
+        opt_level: Kernel optimization level forwarded to cuTile.
+        occupancy: Optional occupancy hint for kernel launch.
+        num_ctas: Optional CTA count hint for kernel launch.
+    Returns:
+        object: Function result value.
+    """
+    ct = _runtime.get_cutile_module()
+    kernel_kwargs: dict[str, Any] = {"opt_level": opt_level}
+    if occupancy is not None:
+        kernel_kwargs["occupancy"] = occupancy
+    if num_ctas is not None:
+        kernel_kwargs["num_ctas"] = num_ctas
+    if len(chunk_plan) != 2:
+        raise ValueError(f"Chunked kernel currently expects exactly 2 chunks, got {chunk_plan!r}.")
+    (off0, w0), (off1, w1) = chunk_plan
+    if accum_mode == "fp16":
+        @ct.kernel(**kernel_kwargs)
+        def flash_fwd_kernel_fp16acc_chunked(q0, k0_t, v0, q1, k1_t, v1, out0, out1, scale):
+            """
+            Run flash fwd kernel fp16acc chunked.
+            It is used during kernel configuration, compilation, or launch.
+            Args:
+                q0: First query chunk for chunked head-dimension kernels.
+                k0_t: Transposed first key chunk.
+                v0: First value chunk for chunked execution.
+                q1: Second query chunk for chunked head-dimension kernels.
+                k1_t: Transposed second key chunk.
+                v1: Second value chunk for chunked execution.
+                out0: First output chunk buffer.
+                out1: Second output chunk buffer.
+                scale: Attention scaling factor.
+            """
+            bh_idx = ct.bid(0)
+            q_tile_idx = ct.bid(1)
+            q0_tile = ct.load(q0, index=(bh_idx, q_tile_idx, 0), shape=(1, tile_m, w0))
+            q1_tile = ct.load(q1, index=(bh_idx, q_tile_idx, 0), shape=(1, tile_m, w1))
+            seq_len = q0.shape[1]
+            num_k_tiles = ct.cdiv(seq_len, tile_n)
+            if causal:
+                causal_cols = ct.minimum((q_tile_idx + 1) * tile_m, seq_len)
+                num_k_tiles = ct.cdiv(causal_cols, tile_n)
+            row = q_tile_idx * tile_m + ct.arange(tile_m, dtype=ct.int32)
+            row = ct.expand_dims(row, 1)
+            row = ct.expand_dims(row, 0)
+            row_in_bounds = row < seq_len
+            m_i = ct.full((1, tile_m, 1), _NEG_LARGE, ct.float32)
+            l_i = ct.zeros((1, tile_m, 1), ct.float32)
+            acc0 = ct.zeros((1, tile_m, w0), dtype)
+            acc1 = ct.zeros((1, tile_m, w1), dtype)
+            for k_tile_idx in range(num_k_tiles):
+                k0_tile_t = ct.load(k0_t, index=(bh_idx, 0, k_tile_idx), shape=(1, w0, tile_n))
+                k1_tile_t = ct.load(k1_t, index=(bh_idx, 0, k_tile_idx), shape=(1, w1, tile_n))
+                score = ct.astype(ct.matmul(q0_tile, k0_tile_t), ct.float32) + ct.astype(
+                    ct.matmul(q1_tile, k1_tile_t), ct.float32
+                )
+                score = score * scale
+                col = k_tile_idx * tile_n + ct.arange(tile_n, dtype=ct.int32)
+                col = ct.expand_dims(col, 0)
+                col = ct.expand_dims(col, 0)
+                key_in_bounds = col < seq_len
+                valid = row_in_bounds & key_in_bounds
+                if causal:
+                    valid = valid & (row >= col)
+                score = ct.where(valid, score, _NEG_LARGE)
+                tile_max = ct.max(score, axis=2, keepdims=True)
+                m_next = ct.maximum(m_i, tile_max)
+                alpha = ct.exp(m_i - m_next)
+                p = ct.exp(score - m_next)
+                l_i = l_i * alpha + ct.sum(p, axis=2, keepdims=True)
+                p_lowp = ct.astype(p, dtype)
+                alpha_lowp = ct.astype(alpha, dtype)
+                v0_tile = ct.load(v0, index=(bh_idx, k_tile_idx, 0), shape=(1, tile_n, w0))
+                v1_tile = ct.load(v1, index=(bh_idx, k_tile_idx, 0), shape=(1, tile_n, w1))
+                acc0 = acc0 * alpha_lowp + ct.matmul(p_lowp, v0_tile)
+                acc1 = acc1 * alpha_lowp + ct.matmul(p_lowp, v1_tile)
+                m_i = m_next
+            safe_l = ct.where(row_in_bounds, l_i, 1.0)
+            inv_l = 1.0 / safe_l
+            inv_l_lowp = ct.astype(inv_l, dtype)
+            out0_tile = ct.where(row_in_bounds, acc0 * inv_l_lowp, 0.0)
+            out1_tile = ct.where(row_in_bounds, acc1 * inv_l_lowp, 0.0)
+            ct.store(out0, index=(bh_idx, q_tile_idx, 0), tile=out0_tile)
+            ct.store(out1, index=(bh_idx, q_tile_idx, 0), tile=out1_tile)
+        return flash_fwd_kernel_fp16acc_chunked
+    @ct.kernel(**kernel_kwargs)
+    def flash_fwd_kernel_fp32acc_chunked(q0, k0_t, v0, q1, k1_t, v1, out0, out1, scale):
+        """
+        Run flash fwd kernel fp32acc chunked.
+        It is used during kernel configuration, compilation, or launch.
+        Args:
+            q0: First query chunk for chunked head-dimension kernels.
+            k0_t: Transposed first key chunk.
+            v0: First value chunk for chunked execution.
+            q1: Second query chunk for chunked head-dimension kernels.
+            k1_t: Transposed second key chunk.
+            v1: Second value chunk for chunked execution.
+            out0: First output chunk buffer.
+            out1: Second output chunk buffer.
+            scale: Attention scaling factor.
+        """
+        bh_idx = ct.bid(0)
+        q_tile_idx = ct.bid(1)
+        q0_tile = ct.load(q0, index=(bh_idx, q_tile_idx, 0), shape=(1, tile_m, w0))
+        q1_tile = ct.load(q1, index=(bh_idx, q_tile_idx, 0), shape=(1, tile_m, w1))
+        seq_len = q0.shape[1]
+        num_k_tiles = ct.cdiv(seq_len, tile_n)
+        if causal:
+            causal_cols = ct.minimum((q_tile_idx + 1) * tile_m, seq_len)
+            num_k_tiles = ct.cdiv(causal_cols, tile_n)
+        row = q_tile_idx * tile_m + ct.arange(tile_m, dtype=ct.int32)
+        row = ct.expand_dims(row, 1)
+        row = ct.expand_dims(row, 0)
+        row_in_bounds = row < seq_len
+        m_i = ct.full((1, tile_m, 1), _NEG_LARGE, ct.float32)
+        l_i = ct.zeros((1, tile_m, 1), ct.float32)
+        acc0 = ct.zeros((1, tile_m, w0), ct.float32)
+        acc1 = ct.zeros((1, tile_m, w1), ct.float32)
+        for k_tile_idx in range(num_k_tiles):
+            k0_tile_t = ct.load(k0_t, index=(bh_idx, 0, k_tile_idx), shape=(1, w0, tile_n))
+            k1_tile_t = ct.load(k1_t, index=(bh_idx, 0, k_tile_idx), shape=(1, w1, tile_n))
+            score = ct.astype(ct.matmul(q0_tile, k0_tile_t), ct.float32) + ct.astype(
+                ct.matmul(q1_tile, k1_tile_t), ct.float32
+            )
+            score = score * scale
+            col = k_tile_idx * tile_n + ct.arange(tile_n, dtype=ct.int32)
+            col = ct.expand_dims(col, 0)
+            col = ct.expand_dims(col, 0)
+            key_in_bounds = col < seq_len
+            valid = row_in_bounds & key_in_bounds
+            if causal:
+                valid = valid & (row >= col)
+            score = ct.where(valid, score, _NEG_LARGE)
+            tile_max = ct.max(score, axis=2, keepdims=True)
+            m_next = ct.maximum(m_i, tile_max)
+            alpha = ct.exp(m_i - m_next)
+            p = ct.exp(score - m_next)
+            l_i = l_i * alpha + ct.sum(p, axis=2, keepdims=True)
+            p_lowp = ct.astype(p, dtype)
+            v0_tile = ct.load(v0, index=(bh_idx, k_tile_idx, 0), shape=(1, tile_n, w0))
+            v1_tile = ct.load(v1, index=(bh_idx, k_tile_idx, 0), shape=(1, tile_n, w1))
+            acc0 = acc0 * alpha + ct.astype(ct.matmul(p_lowp, v0_tile), ct.float32)
+            acc1 = acc1 * alpha + ct.astype(ct.matmul(p_lowp, v1_tile), ct.float32)
+            m_i = m_next
+        safe_l = ct.where(row_in_bounds, l_i, 1.0)
+        inv_l = 1.0 / safe_l
+        out0_tile = ct.astype(ct.where(row_in_bounds, acc0 * inv_l, 0.0), dtype)
+        out1_tile = ct.astype(ct.where(row_in_bounds, acc1 * inv_l, 0.0), dtype)
+        ct.store(out0, index=(bh_idx, q_tile_idx, 0), tile=out0_tile)
+        ct.store(out1, index=(bh_idx, q_tile_idx, 0), tile=out1_tile)
+    return flash_fwd_kernel_fp32acc_chunked
+def _launch_cutile_kernel(
+    ct: Any,
+    cupy_mod: Any,
+    kernel: Any,
+    grid: tuple[int, int, int],
+    args: tuple[Any, ...],
+) -> None:
+    """
+    Internal helper for launch cutile kernel.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        ct: Imported cuTile module instance.
+        cupy_mod: Imported CuPy module instance.
+        kernel: Compiled kernel object to launch.
+        grid: Function argument.
+        args: Parsed command-line arguments namespace.
+    """
+    torch_mod = _runtime.get_torch_module()
+    sync_mode = os.getenv("TILEDATTN_SYNC_MODE", "async").strip().lower()
+    stream = cupy_mod.cuda.get_current_stream()
+    if sync_mode == "strict":
+        torch_mod.cuda.synchronize()
+    ct.launch(stream, grid, kernel, args)
+    if sync_mode in {"strict", "post"}:
+        stream.synchronize()
+    elif sync_mode == "async":
+        pass
+    else:
+        raise ValueError(
+            f"Unsupported TILEDATTN_SYNC_MODE={sync_mode!r}. "
+            "Use one of: strict, post, async."
+        )
+def run_flash_fwd(
+    q,
+    k,
+    v,
+    *,
+    causal: bool,
+    scale: float,
+):
+    """
+    Run flash fwd.
+    It is used during kernel configuration, compilation, or launch.
+    Args:
+        q: Query tensor in attention layout.
+        k: Key tensor in attention layout.
+        v: Value tensor in attention layout.
+        causal: Whether causal masking is enabled.
+        scale: Attention scaling factor.
+    Returns:
+        object: Function result value.
+    """
+    torch_mod = _runtime.get_torch_module()
+    cupy_mod = _runtime.get_cupy_module()
+    ct = _runtime.get_cutile_module()
+    opt_level, occupancy, num_ctas = _resolve_kernel_options()
+    batch, heads, seq_len, head_dim = map(int, q.shape)
+    accum_mode = _resolve_accum_mode(seq_len=seq_len, head_dim=head_dim, causal=causal)
+    bh = batch * heads
+    if os.getenv("TILEDATTN_TILE_M") or os.getenv("TILEDATTN_TILE_N"):
+        tile_m, tile_n = _resolve_tile_config()
+    else:
+        tile_m, tile_n = _default_tile_config_for_shape(
+            seq_len=seq_len, head_dim=head_dim, causal=causal
+        )
+    chunk_plan = _resolve_chunk_plan(head_dim)
+    if chunk_plan is not None:
+        kernel_head_dim = head_dim
+        pad_dim = 0
+    else:
+        kernel_head_dim, pad_dim = _resolve_kernel_head_dim(head_dim)
+    q_bh = q.contiguous().reshape(bh, seq_len, head_dim)
+    k_bh = k.contiguous().reshape(bh, seq_len, head_dim)
+    v_bh = v.contiguous().reshape(bh, seq_len, head_dim)
+    if pad_dim > 0:
+        pad = (0, pad_dim)
+        q_bh = torch_mod.nn.functional.pad(q_bh, pad)
+        k_bh = torch_mod.nn.functional.pad(k_bh, pad)
+        v_bh = torch_mod.nn.functional.pad(v_bh, pad)
+    # Keep K as a transpose view to avoid a per-call materialization copy.
+    k_t = k_bh.transpose(1, 2)
+    out_bh = torch_mod.empty_like(q_bh)
+    out_ct_dtype = _ct_dtype_for_torch_dtype(ct, q.dtype)
+    use_aligned_fastpath = _should_use_aligned_noncausal_fastpath(
+        seq_len=seq_len,
+        head_dim=head_dim,
+        causal=causal,
+        pad_dim=pad_dim,
+        chunk_plan=chunk_plan,
+        tile_m=tile_m,
+        tile_n=tile_n,
+    )
+    if chunk_plan is not None:
+        kernel_variant = "chunked"
+    elif use_aligned_fastpath:
+        kernel_variant = "direct_aligned_noncausal"
+    else:
+        kernel_variant = "direct"
+    kernel_key = (
+        tile_m,
+        tile_n,
+        kernel_head_dim,
+        kernel_variant,
+        str(q.dtype),
+        causal,
+        accum_mode,
+        opt_level,
+        occupancy if occupancy is not None else -1,
+        num_ctas if num_ctas is not None else -1,
+    )
+    if chunk_plan is None:
+        if use_aligned_fastpath:
+            kernel = get_kernel(
+                kernel_key,
+                lambda: make_flashattn_fwd_kernel_aligned_noncausal(
+                    tile_m,
+                    tile_n,
+                    kernel_head_dim,
+                    dtype=out_ct_dtype,
+                    accum_mode=accum_mode,
+                    opt_level=opt_level,
+                    occupancy=occupancy,
+                    num_ctas=num_ctas,
+                ),
+            )
+        else:
+            kernel = get_kernel(
+                kernel_key,
+                lambda: make_flashattn_fwd_kernel(
+                    tile_m,
+                    tile_n,
+                    kernel_head_dim,
+                    dtype=out_ct_dtype,
+                    causal=causal,
+                    accum_mode=accum_mode,
+                    opt_level=opt_level,
+                    occupancy=occupancy,
+                    num_ctas=num_ctas,
+                ),
+            )
+    else:
+        kernel = get_kernel(
+            kernel_key,
+            lambda: make_flashattn_fwd_kernel_chunked(
+                tile_m,
+                tile_n,
+                kernel_head_dim,
+                chunk_plan=chunk_plan,
+                dtype=out_ct_dtype,
+                causal=causal,
+                accum_mode=accum_mode,
+                opt_level=opt_level,
+                occupancy=occupancy,
+                num_ctas=num_ctas,
+            ),
+        )
+    grid = (bh, (seq_len + tile_m - 1) // tile_m, 1)
+    if chunk_plan is None:
+        args = (q_bh, k_t, v_bh, out_bh, float(scale))
+    else:
+        (off0, w0), (off1, w1) = chunk_plan
+        if off0 != 0 or off1 != w0 or (w0 + w1) != head_dim:
+            raise ValueError(f"Unexpected chunk plan for head_dim={head_dim}: {chunk_plan!r}")
+        q0 = q_bh[:, :, :w0]
+        q1 = q_bh[:, :, w0:]
+        k0_t = k_bh[:, :, :w0].transpose(1, 2)
+        k1_t = k_bh[:, :, w0:].transpose(1, 2)
+        v0 = v_bh[:, :, :w0]
+        v1 = v_bh[:, :, w0:]
+        out0 = out_bh[:, :, :w0]
+        out1 = out_bh[:, :, w0:]
+        args = (q0, k0_t, v0, q1, k1_t, v1, out0, out1, float(scale))
+    _launch_cutile_kernel(ct, cupy_mod, kernel, grid, args)
+    out = out_bh[:, :, :head_dim] if pad_dim > 0 else out_bh
+    return out.reshape(batch, heads, seq_len, head_dim)

build/torch-cuda/metadata.json ADDED Viewed

	@@ -0,0 +1,10 @@

+{
+  "name": "tiledattention",
+  "id": "_tiledattention_cuda_703d09c_dirty",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda"
+  }
+}

build/torch-cuda/sdpa.py ADDED Viewed

	@@ -0,0 +1,54 @@

+"""PyTorch-facing SDPA API."""
+from __future__ import annotations
+import math
+from numbers import Real
+try:
+    from . import _runtime
+    from ._errors import InvalidShapeError
+    from .kernels import run_flash_fwd
+    from .utils.checks import validate_sdpa_inputs
+except ImportError:  # pragma: no cover - flattened kernel package layout
+    import _runtime  # type: ignore
+    from _errors import InvalidShapeError  # type: ignore
+    from kernels import run_flash_fwd  # type: ignore
+    from utils.checks import validate_sdpa_inputs  # type: ignore
+def sdpa(
+    q,
+    k,
+    v,
+    *,
+    causal: bool = False,
+    scale: float | None = None,
+):
+    """
+    Compute scaled dot-product attention.
+    It is part of the public SDPA execution path.
+    Args:
+        q: Query tensor in attention layout.
+        k: Key tensor in attention layout.
+        v: Value tensor in attention layout.
+        causal: Whether causal masking is enabled.
+        scale: Attention scaling factor.
+    Returns:
+        object: Function result value.
+    """
+    if not isinstance(causal, bool):
+        raise InvalidShapeError("causal must be a bool.")
+    if scale is not None and (not isinstance(scale, Real) or float(scale) <= 0):
+        raise InvalidShapeError("scale must be a positive number when provided.")
+    _runtime.require_supported_runtime()
+    torch_mod = _runtime.get_torch_module()
+    validate_sdpa_inputs(torch_mod, q, k, v)
+    head_dim = int(q.shape[-1])
+    resolved_scale = float(scale) if scale is not None else 1.0 / math.sqrt(head_dim)
+    return run_flash_fwd(q, k, v, causal=causal, scale=resolved_scale)

build/torch-cuda/tiledattention/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch-cuda/utils/__init__.py ADDED Viewed

	@@ -0,0 +1,7 @@

+"""Utility helpers for tiledattention."""
+from __future__ import annotations
+from .checks import validate_sdpa_inputs
+__all__ = ["validate_sdpa_inputs"]

build/torch-cuda/utils/checks.py ADDED Viewed

	@@ -0,0 +1,58 @@

+"""Input checks shared by tiledattention APIs."""
+from __future__ import annotations
+from types import ModuleType
+try:
+    from .._errors import DTypeNotSupportedError, InvalidShapeError
+except ImportError:  # pragma: no cover - flattened kernel package layout
+    from _errors import DTypeNotSupportedError, InvalidShapeError  # type: ignore
+def _ensure_tensor(torch_mod: ModuleType, name: str, value: object) -> None:
+    """
+    Internal helper for ensure tensor.
+    It enforces input constraints before kernel execution.
+    Args:
+        torch_mod: Imported torch module instance.
+        name: Identifier or metric name.
+        value: Scalar value to parse or format.
+    """
+    if not isinstance(value, torch_mod.Tensor):
+        raise InvalidShapeError(f"{name} must be a torch.Tensor.")
+def validate_sdpa_inputs(torch_mod: ModuleType, q, k, v) -> None:
+    """
+    Validate sdpa inputs.
+    It enforces input constraints before kernel execution.
+    Args:
+        torch_mod: Imported torch module instance.
+        q: Query tensor in attention layout.
+        k: Key tensor in attention layout.
+        v: Value tensor in attention layout.
+    """
+    _ensure_tensor(torch_mod, "q", q)
+    _ensure_tensor(torch_mod, "k", k)
+    _ensure_tensor(torch_mod, "v", v)
+    for name, tensor in (("q", q), ("k", k), ("v", v)):
+        if tensor.ndim != 4:
+            raise InvalidShapeError(f"{name} must have shape [B, H, S, D].")
+        if not bool(tensor.is_cuda):
+            raise InvalidShapeError(f"{name} must be a CUDA tensor.")
+    if q.shape != k.shape or q.shape != v.shape:
+        raise InvalidShapeError("q, k, and v must have matching shape [B, H, S, D].")
+    if q.device != k.device or q.device != v.device:
+        raise InvalidShapeError("q, k, and v must be on the same CUDA device.")
+    allowed_dtypes = {torch_mod.float16, torch_mod.bfloat16}
+    if q.dtype not in allowed_dtypes:
+        raise DTypeNotSupportedError("q dtype must be torch.float16 or torch.bfloat16.")
+    if q.dtype != k.dtype or q.dtype != v.dtype:
+        raise DTypeNotSupportedError("q, k, and v must have identical dtype.")