Duplicate from kernels-community/metal-flash-sdpa

Co-authored-by: Eric Buehler <EricB@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text
+*.metallib filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,27 @@
+---
+library_name: kernels
+license: apache-2.0
+---
+
+This is the repository card of kernels-community/metal-flash-sdpa that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+kernel_module = get_kernel("kernels-community/metal-flash-sdpa")
+flash_attention_varlen = kernel_module.flash_attention_varlen
+
+flash_attention_varlen(...)
+```
+
+## Available functions
+- `flash_attention_varlen`
+- `flash_attn_varlen_func`
+- `ops`
+
+## Benchmarks
+
+No benchmark available yet.

benchmark.py

@@ -0,0 +1,304 @@
+#!/usr/bin/env python3
+"""Benchmark causal mask performance scaling with sequence length"""
+
+import torch
+import time
+import matplotlib.pyplot as plt
+import numpy as np
+from typing import List
+import kernels
+
+metal_flash_sdpa = kernels.get_kernel("kernels-community/metal-flash-sdpa")
+
+
+def create_cu_seqlens(seq_lengths: List[int]) -> torch.Tensor:
+    """Create cumulative sequence lengths tensor."""
+    cu_seqlens = [0]
+    for length in seq_lengths:
+        cu_seqlens.append(cu_seqlens[-1] + length)
+    return torch.tensor(cu_seqlens, dtype=torch.int32, device="mps")
+
+
+def benchmark_flash_sdpa_causal(
+    batch_size: int,
+    num_heads: int,
+    seq_len: int,
+    head_dim: int,
+    dtype: torch.dtype,
+    num_iterations: int = 20,
+) -> float:
+    """Benchmark Flash SDPA with causal mask"""
+
+    seq_lengths = [seq_len] * batch_size
+    cu_seqlens = create_cu_seqlens(seq_lengths)
+    total_tokens = sum(seq_lengths)
+
+    # Create input tensors
+    query = torch.randn(total_tokens, num_heads, head_dim, dtype=dtype, device="mps")
+    key = torch.randn(total_tokens, num_heads, head_dim, dtype=dtype, device="mps")
+    value = torch.randn(total_tokens, num_heads, head_dim, dtype=dtype, device="mps")
+    out = torch.empty_like(query)
+
+    scale = 1.0 / (head_dim**0.5)
+
+    # Warmup
+    for _ in range(5):
+        metal_flash_sdpa.flash_attention_varlen(
+            out=out,
+            query=query,
+            key=key,
+            value=value,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=seq_len,
+            max_seqlen_k=seq_len,
+            do_causal=True,
+            scale=scale,
+            softcapping=1.0,
+        )
+    torch.mps.synchronize()
+
+    # Benchmark
+    start_time = time.perf_counter()
+    for _ in range(num_iterations):
+        metal_flash_sdpa.flash_attention_varlen(
+            out=out,
+            query=query,
+            key=key,
+            value=value,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=seq_len,
+            max_seqlen_k=seq_len,
+            do_causal=True,
+            scale=scale,
+            softcapping=1.0,
+        )
+    torch.mps.synchronize()
+    end_time = time.perf_counter()
+
+    return (end_time - start_time) * 1000 / num_iterations
+
+
+def benchmark_naive_sdpa_causal(
+    batch_size: int,
+    num_heads: int,
+    seq_len: int,
+    head_dim: int,
+    dtype: torch.dtype,
+    num_iterations: int = 20,
+) -> float:
+    """Benchmark naive SDPA with causal mask"""
+
+    # Create input tensors
+    query = torch.randn(
+        batch_size, num_heads, seq_len, head_dim, dtype=dtype, device="mps"
+    )
+    key = torch.randn(
+        batch_size, num_heads, seq_len, head_dim, dtype=dtype, device="mps"
+    )
+    value = torch.randn(
+        batch_size, num_heads, seq_len, head_dim, dtype=dtype, device="mps"
+    )
+
+    scale = 1.0 / (head_dim**0.5)
+
+    # Precompute causal mask
+    mask = torch.triu(torch.ones(seq_len, seq_len, device="mps"), diagonal=1).bool()
+
+    # Warmup
+    for _ in range(5):
+        scores = torch.matmul(query, key.transpose(-2, -1)) * scale
+        scores = scores.masked_fill(mask, float("-inf"))
+        attn_weights = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn_weights, value)
+    torch.mps.synchronize()
+
+    # Benchmark
+    start_time = time.perf_counter()
+    for _ in range(num_iterations):
+        scores = torch.matmul(query, key.transpose(-2, -1)) * scale
+        scores = scores.masked_fill(mask, float("-inf"))
+        attn_weights = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn_weights, value)
+    torch.mps.synchronize()
+    end_time = time.perf_counter()
+
+    return (end_time - start_time) * 1000 / num_iterations
+
+
+def run_scaling_benchmark():
+    """Run causal mask scaling benchmark"""
+
+    print("=" * 80)
+    print("Causal Mask Performance Scaling Benchmark")
+    print("Batch Size: 4, Head Dimension: 64")
+    print("=" * 80)
+
+    # Configuration
+    batch_size = 4
+    num_heads = 16
+    head_dim = 64
+    dtype = torch.float16
+
+    # Sequence lengths from 512 to 4096
+    seq_lengths = [512, 768, 1024, 1536, 2048, 3072, 4096]
+
+    flash_times = []
+    naive_times = []
+    speedups = []
+
+    print(f"{'Seq Len':<8} {'Flash (ms)':<12} {'Naive (ms)':<12} {'Speedup':<10}")
+    print("-" * 50)
+
+    for seq_len in seq_lengths:
+        # Benchmark Flash SDPA
+        flash_time = benchmark_flash_sdpa_causal(
+            batch_size, num_heads, seq_len, head_dim, dtype
+        )
+        flash_times.append(flash_time)
+
+        # Benchmark Naive SDPA
+        naive_time = benchmark_naive_sdpa_causal(
+            batch_size, num_heads, seq_len, head_dim, dtype
+        )
+        naive_times.append(naive_time)
+
+        speedup = naive_time / flash_time
+        speedups.append(speedup)
+
+        print(f"{seq_len:<8} {flash_time:<12.2f} {naive_time:<12.2f} {speedup:<10.2f}x")
+
+    return seq_lengths, flash_times, naive_times, speedups
+
+
+def create_line_plot(seq_lengths, flash_times, naive_times, speedups):
+    """Create line graph visualization"""
+
+    # Create figure with single plot
+    fig, ax = plt.subplots(1, 1, figsize=(12, 8))
+    fig.suptitle(
+        "Causal Mask Performance Scaling\n(Batch Size: 4, Head Dimension: 64)",
+        fontsize=16,
+    )
+
+    # Plot execution times
+    ax.plot(
+        seq_lengths,
+        flash_times,
+        marker="o",
+        linewidth=3,
+        markersize=10,
+        label="Flash SDPA",
+        color="blue",
+    )
+    ax.plot(
+        seq_lengths,
+        naive_times,
+        marker="s",
+        linewidth=3,
+        markersize=10,
+        label="Naive SDPA",
+        color="red",
+    )
+
+    ax.set_xlabel("Sequence Length", fontsize=14)
+    ax.set_ylabel("Time (ms)", fontsize=14)
+    ax.set_title("Execution Time vs Sequence Length", fontsize=16)
+    ax.grid(True, alpha=0.3)
+    ax.legend(fontsize=12)
+
+    # Add value annotations for all points
+    for i, (seq_len, flash_time, naive_time) in enumerate(
+        zip(seq_lengths, flash_times, naive_times)
+    ):
+        ax.annotate(
+            f"{flash_time:.1f}ms",
+            xy=(seq_len, flash_time),
+            xytext=(5, 5),
+            textcoords="offset points",
+            fontsize=10,
+            color="blue",
+        )
+        ax.annotate(
+            f"{naive_time:.1f}ms",
+            xy=(seq_len, naive_time),
+            xytext=(5, 5),
+            textcoords="offset points",
+            fontsize=10,
+            color="red",
+        )
+
+    # Set axis limits to better show the data
+    ax.set_xlim(seq_lengths[0] - 100, seq_lengths[-1] + 100)
+    ax.set_ylim(0, max(naive_times) * 1.1)
+
+    plt.tight_layout()
+    plt.savefig("benchmark.png", dpi=300, bbox_inches="tight")
+    plt.show()
+
+
+def print_analysis(seq_lengths, flash_times, naive_times, speedups):
+    """Print detailed analysis of the results"""
+
+    print("\n" + "=" * 80)
+    print("DETAILED ANALYSIS")
+    print("=" * 80)
+
+    # Performance scaling analysis
+    print("\n1. Performance Scaling:")
+    print(
+        f"   • Flash SDPA: {flash_times[0]:.2f}ms → {flash_times[-1]:.2f}ms ({flash_times[-1] / flash_times[0]:.1f}x increase)"
+    )
+    print(
+        f"   • Naive SDPA: {naive_times[0]:.2f}ms → {naive_times[-1]:.2f}ms ({naive_times[-1] / naive_times[0]:.1f}x increase)"
+    )
+
+    # Speedup analysis
+    print("\n2. Speedup Analysis:")
+    print(f"   • Average Speedup: {np.mean(speedups):.2f}x")
+    print(
+        f"   • Max Speedup: {np.max(speedups):.2f}x (at seq_len={seq_lengths[np.argmax(speedups)]})"
+    )
+    print(
+        f"   • Min Speedup: {np.min(speedups):.2f}x (at seq_len={seq_lengths[np.argmin(speedups)]})"
+    )
+
+    # Efficiency analysis
+    print("\n3. Efficiency Analysis:")
+    speedup_improvement = speedups[-1] / speedups[0]
+    print(f"   • Speedup improvement from 512→4096: {speedup_improvement:.2f}x")
+
+    if speedup_improvement > 1.1:
+        print("   • Flash SDPA becomes MORE efficient at longer sequences")
+    elif speedup_improvement < 0.9:
+        print("   • Flash SDPA becomes LESS efficient at longer sequences")
+    else:
+        print("   • Flash SDPA maintains consistent efficiency across sequence lengths")
+
+    # Memory complexity analysis
+    print("\n4. Theoretical Complexity:")
+    print(f"   • Sequence length increased by: {seq_lengths[-1] / seq_lengths[0]:.1f}x")
+    print(
+        f"   • Theoretical O(n²) complexity increase: {(seq_lengths[-1] / seq_lengths[0]) ** 2:.1f}x"
+    )
+    print(f"   • Actual Flash SDPA increase: {flash_times[-1] / flash_times[0]:.1f}x")
+    efficiency_ratio = (flash_times[-1] / flash_times[0]) / (
+        (seq_lengths[-1] / seq_lengths[0]) ** 2
+    )
+    print(f"   • Flash SDPA efficiency ratio: {efficiency_ratio:.3f} (lower is better)")
+
+
+def main():
+    # Run the scaling benchmark
+    seq_lengths, flash_times, naive_times, speedups = run_scaling_benchmark()
+
+    # Create line plot visualization
+    create_line_plot(seq_lengths, flash_times, naive_times, speedups)
+
+    # Print detailed analysis
+    print_analysis(seq_lengths, flash_times, naive_times, speedups)
+
+
+if __name__ == "__main__":
+    main()

build.toml

@@ -0,0 +1,18 @@
+[general]
+name = "metal_flash_sdpa"
+universal = false
+
+[torch]
+src = [
+    "torch-ext/torch_binding.cpp",
+    "torch-ext/torch_binding.h",
+]
+
+[kernel.sdpa_metal]
+backend = "metal"
+src = [
+  "sdpa-metal/scaled_dot_product_attention.mm",
+  "sdpa-metal/scaled_dot_product_attention.metal",
+  "sdpa-metal/common.h",
+]
+depends = [ "torch" ]

build/torch210-metal-aarch64-darwin/__init__.py

@@ -0,0 +1,11 @@
+from ._custom_ops import (
+    flash_attention_varlen,
+    flash_attn_varlen_func,
+)
+from ._ops import ops
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+    "ops",
+]

build/torch210-metal-aarch64-darwin/_custom_ops.py

@@ -0,0 +1,117 @@
+from typing import List, Optional
+
+import torch
+
+from ._ops import ops
+
+
+def flash_attention_varlen(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    do_causal: bool = False,
+    scale: Optional[float] = None,
+    softcapping: float = 1.0,
+) -> None:
+    """
+    Flash Attention with variable-length sequences.
+    
+    Args:
+        out: Output tensor of shape [total_q_tokens, num_heads, head_dim]
+        query: Query tensor of shape [total_q_tokens, num_heads, head_dim]
+        key: Key tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        value: Value tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        cu_seqlens_q: Cumulative sequence lengths for queries, shape [batch_size + 1], dtype must be torch.int32
+        cu_seqlens_k: Cumulative sequence lengths for keys, shape [batch_size + 1], dtype must be torch.int32
+        max_seqlen_q: Maximum sequence length in the query batch
+        max_seqlen_k: Maximum sequence length in the key batch
+        do_causal: Whether to apply causal masking
+        scale: Attention scale factor (default: 1/sqrt(head_dim))
+        softcapping: Softcapping value (default: 1.0, must be 1.0 for this implementation)
+    
+    Note:
+        - cu_seqlens_q and cu_seqlens_k must have dtype torch.int32 for Metal compatibility
+        - Supported head dimensions: 32, 64, 72, 80, 96, 128
+        - Masks are not supported
+    """
+    if scale is None:
+        scale = query.shape[-1] ** -0.5
+    
+    ops.flash_attention_varlen(
+        out,
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        do_causal,
+        scale,
+        softcapping,
+    )
+
+def flash_attn_varlen_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: tuple = (-1, -1),
+    alibi_slopes: Optional[torch.Tensor] = None,
+    deterministic: bool = False,
+    return_attn_probs: bool = False,
+) -> torch.Tensor:
+    """
+    Flash Attention function with API compatible with the original Flash Attention.
+    
+    Note: This implementation does not support:
+    - dropout
+    - window attention
+    - alibi slopes
+    - returning attention probabilities
+    """
+    if dropout_p > 0:
+        raise NotImplementedError("Dropout is not supported in this implementation")
+    if window_size != (-1, -1):
+        raise NotImplementedError("Window attention is not supported")
+    if alibi_slopes is not None:
+        raise NotImplementedError("ALiBi is not supported")
+    if return_attn_probs:
+        raise NotImplementedError("Returning attention probabilities is not supported")
+    
+    # Create output tensor
+    out = torch.empty_like(q)
+    
+    # Call the kernel
+    flash_attention_varlen(
+        out=out,
+        query=q,
+        key=k,
+        value=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        do_causal=causal,
+        scale=softmax_scale,
+        softcapping=1.0,
+    )
+    
+    return out
+
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+]

build/torch210-metal-aarch64-darwin/_metal_flash_sdpa_metal_9cd92d8.abi3.so

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

build/torch210-metal-aarch64-darwin/_ops.py

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

build/torch210-metal-aarch64-darwin/metadata.json

@@ -0,0 +1,10 @@
+{
+  "name": "metal-flash-sdpa",
+  "id": "_metal_flash_sdpa_metal_9cd92d8",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "metal"
+  }
+}

build/torch210-metal-aarch64-darwin/metal_flash_sdpa/__init__.py

@@ -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/torch211-metal-aarch64-darwin/__init__.py

@@ -0,0 +1,11 @@
+from ._custom_ops import (
+    flash_attention_varlen,
+    flash_attn_varlen_func,
+)
+from ._ops import ops
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+    "ops",
+]

build/torch211-metal-aarch64-darwin/_custom_ops.py

@@ -0,0 +1,117 @@
+from typing import List, Optional
+
+import torch
+
+from ._ops import ops
+
+
+def flash_attention_varlen(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    do_causal: bool = False,
+    scale: Optional[float] = None,
+    softcapping: float = 1.0,
+) -> None:
+    """
+    Flash Attention with variable-length sequences.
+    
+    Args:
+        out: Output tensor of shape [total_q_tokens, num_heads, head_dim]
+        query: Query tensor of shape [total_q_tokens, num_heads, head_dim]
+        key: Key tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        value: Value tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        cu_seqlens_q: Cumulative sequence lengths for queries, shape [batch_size + 1], dtype must be torch.int32
+        cu_seqlens_k: Cumulative sequence lengths for keys, shape [batch_size + 1], dtype must be torch.int32
+        max_seqlen_q: Maximum sequence length in the query batch
+        max_seqlen_k: Maximum sequence length in the key batch
+        do_causal: Whether to apply causal masking
+        scale: Attention scale factor (default: 1/sqrt(head_dim))
+        softcapping: Softcapping value (default: 1.0, must be 1.0 for this implementation)
+    
+    Note:
+        - cu_seqlens_q and cu_seqlens_k must have dtype torch.int32 for Metal compatibility
+        - Supported head dimensions: 32, 64, 72, 80, 96, 128
+        - Masks are not supported
+    """
+    if scale is None:
+        scale = query.shape[-1] ** -0.5
+    
+    ops.flash_attention_varlen(
+        out,
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        do_causal,
+        scale,
+        softcapping,
+    )
+
+def flash_attn_varlen_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: tuple = (-1, -1),
+    alibi_slopes: Optional[torch.Tensor] = None,
+    deterministic: bool = False,
+    return_attn_probs: bool = False,
+) -> torch.Tensor:
+    """
+    Flash Attention function with API compatible with the original Flash Attention.
+    
+    Note: This implementation does not support:
+    - dropout
+    - window attention
+    - alibi slopes
+    - returning attention probabilities
+    """
+    if dropout_p > 0:
+        raise NotImplementedError("Dropout is not supported in this implementation")
+    if window_size != (-1, -1):
+        raise NotImplementedError("Window attention is not supported")
+    if alibi_slopes is not None:
+        raise NotImplementedError("ALiBi is not supported")
+    if return_attn_probs:
+        raise NotImplementedError("Returning attention probabilities is not supported")
+    
+    # Create output tensor
+    out = torch.empty_like(q)
+    
+    # Call the kernel
+    flash_attention_varlen(
+        out=out,
+        query=q,
+        key=k,
+        value=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        do_causal=causal,
+        scale=softmax_scale,
+        softcapping=1.0,
+    )
+    
+    return out
+
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+]

build/torch211-metal-aarch64-darwin/_metal_flash_sdpa_metal_9cd92d8.abi3.so

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

build/torch211-metal-aarch64-darwin/_ops.py

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

build/torch211-metal-aarch64-darwin/metadata.json

@@ -0,0 +1,10 @@
+{
+  "name": "metal-flash-sdpa",
+  "id": "_metal_flash_sdpa_metal_9cd92d8",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "metal"
+  }
+}

build/torch211-metal-aarch64-darwin/metal_flash_sdpa/__init__.py

@@ -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/torch212-metal-aarch64-darwin/__init__.py

@@ -0,0 +1,11 @@
+from ._custom_ops import (
+    flash_attention_varlen,
+    flash_attn_varlen_func,
+)
+from ._ops import ops
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+    "ops",
+]

build/torch212-metal-aarch64-darwin/_custom_ops.py

@@ -0,0 +1,117 @@
+from typing import List, Optional
+
+import torch
+
+from ._ops import ops
+
+
+def flash_attention_varlen(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    do_causal: bool = False,
+    scale: Optional[float] = None,
+    softcapping: float = 1.0,
+) -> None:
+    """
+    Flash Attention with variable-length sequences.
+    
+    Args:
+        out: Output tensor of shape [total_q_tokens, num_heads, head_dim]
+        query: Query tensor of shape [total_q_tokens, num_heads, head_dim]
+        key: Key tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        value: Value tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        cu_seqlens_q: Cumulative sequence lengths for queries, shape [batch_size + 1], dtype must be torch.int32
+        cu_seqlens_k: Cumulative sequence lengths for keys, shape [batch_size + 1], dtype must be torch.int32
+        max_seqlen_q: Maximum sequence length in the query batch
+        max_seqlen_k: Maximum sequence length in the key batch
+        do_causal: Whether to apply causal masking
+        scale: Attention scale factor (default: 1/sqrt(head_dim))
+        softcapping: Softcapping value (default: 1.0, must be 1.0 for this implementation)
+    
+    Note:
+        - cu_seqlens_q and cu_seqlens_k must have dtype torch.int32 for Metal compatibility
+        - Supported head dimensions: 32, 64, 72, 80, 96, 128
+        - Masks are not supported
+    """
+    if scale is None:
+        scale = query.shape[-1] ** -0.5
+    
+    ops.flash_attention_varlen(
+        out,
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        do_causal,
+        scale,
+        softcapping,
+    )
+
+def flash_attn_varlen_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: tuple = (-1, -1),
+    alibi_slopes: Optional[torch.Tensor] = None,
+    deterministic: bool = False,
+    return_attn_probs: bool = False,
+) -> torch.Tensor:
+    """
+    Flash Attention function with API compatible with the original Flash Attention.
+    
+    Note: This implementation does not support:
+    - dropout
+    - window attention
+    - alibi slopes
+    - returning attention probabilities
+    """
+    if dropout_p > 0:
+        raise NotImplementedError("Dropout is not supported in this implementation")
+    if window_size != (-1, -1):
+        raise NotImplementedError("Window attention is not supported")
+    if alibi_slopes is not None:
+        raise NotImplementedError("ALiBi is not supported")
+    if return_attn_probs:
+        raise NotImplementedError("Returning attention probabilities is not supported")
+    
+    # Create output tensor
+    out = torch.empty_like(q)
+    
+    # Call the kernel
+    flash_attention_varlen(
+        out=out,
+        query=q,
+        key=k,
+        value=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        do_causal=causal,
+        scale=softmax_scale,
+        softcapping=1.0,
+    )
+    
+    return out
+
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+]

build/torch212-metal-aarch64-darwin/_metal_flash_sdpa_metal_9cd92d8.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8cc31f8453fb18e703bb77d1ea793c8039d14156078eda8c20713af96a7e2e9c
+size 733128

build/torch212-metal-aarch64-darwin/_ops.py

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

build/torch212-metal-aarch64-darwin/metadata.json

@@ -0,0 +1,10 @@
+{
+  "name": "metal-flash-sdpa",
+  "id": "_metal_flash_sdpa_metal_9cd92d8",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "metal"
+  }
+}

build/torch212-metal-aarch64-darwin/metal_flash_sdpa/__init__.py

@@ -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/torch27-metal-aarch64-darwin/metal_flash_sdpa/__init__.py

@@ -0,0 +1,11 @@
+from ._custom_ops import (
+    flash_attention_varlen,
+    flash_attn_varlen_func,
+)
+from ._ops import ops
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+    "ops",
+]

build/torch27-metal-aarch64-darwin/metal_flash_sdpa/_custom_ops.py

@@ -0,0 +1,117 @@
+from typing import List, Optional
+
+import torch
+
+from ._ops import ops
+
+
+def flash_attention_varlen(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    do_causal: bool = False,
+    scale: Optional[float] = None,
+    softcapping: float = 1.0,
+) -> None:
+    """
+    Flash Attention with variable-length sequences.
+    
+    Args:
+        out: Output tensor of shape [total_q_tokens, num_heads, head_dim]
+        query: Query tensor of shape [total_q_tokens, num_heads, head_dim]
+        key: Key tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        value: Value tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        cu_seqlens_q: Cumulative sequence lengths for queries, shape [batch_size + 1], dtype must be torch.int32
+        cu_seqlens_k: Cumulative sequence lengths for keys, shape [batch_size + 1], dtype must be torch.int32
+        max_seqlen_q: Maximum sequence length in the query batch
+        max_seqlen_k: Maximum sequence length in the key batch
+        do_causal: Whether to apply causal masking
+        scale: Attention scale factor (default: 1/sqrt(head_dim))
+        softcapping: Softcapping value (default: 1.0, must be 1.0 for this implementation)
+    
+    Note:
+        - cu_seqlens_q and cu_seqlens_k must have dtype torch.int32 for Metal compatibility
+        - Supported head dimensions: 32, 64, 72, 80, 96, 128
+        - Masks are not supported
+    """
+    if scale is None:
+        scale = query.shape[-1] ** -0.5
+    
+    ops.flash_attention_varlen(
+        out,
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        do_causal,
+        scale,
+        softcapping,
+    )
+
+def flash_attn_varlen_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: tuple = (-1, -1),
+    alibi_slopes: Optional[torch.Tensor] = None,
+    deterministic: bool = False,
+    return_attn_probs: bool = False,
+) -> torch.Tensor:
+    """
+    Flash Attention function with API compatible with the original Flash Attention.
+    
+    Note: This implementation does not support:
+    - dropout
+    - window attention
+    - alibi slopes
+    - returning attention probabilities
+    """
+    if dropout_p > 0:
+        raise NotImplementedError("Dropout is not supported in this implementation")
+    if window_size != (-1, -1):
+        raise NotImplementedError("Window attention is not supported")
+    if alibi_slopes is not None:
+        raise NotImplementedError("ALiBi is not supported")
+    if return_attn_probs:
+        raise NotImplementedError("Returning attention probabilities is not supported")
+    
+    # Create output tensor
+    out = torch.empty_like(q)
+    
+    # Call the kernel
+    flash_attention_varlen(
+        out=out,
+        query=q,
+        key=k,
+        value=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        do_causal=causal,
+        scale=softmax_scale,
+        softcapping=1.0,
+    )
+    
+    return out
+
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+]

build/torch27-metal-aarch64-darwin/metal_flash_sdpa/_metal_flash_sdpa_8aa66c2.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29aa4a03d9cbe5d1b56aebaf0eff43eba5dc0a5796909c2ede503e9e1f80c565
+size 732048

build/torch27-metal-aarch64-darwin/metal_flash_sdpa/_ops.py

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

build/torch29-metal-aarch64-darwin/__init__.py

@@ -0,0 +1,11 @@
+from ._custom_ops import (
+    flash_attention_varlen,
+    flash_attn_varlen_func,
+)
+from ._ops import ops
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+    "ops",
+]

build/torch29-metal-aarch64-darwin/_custom_ops.py

@@ -0,0 +1,117 @@
+from typing import List, Optional
+
+import torch
+
+from ._ops import ops
+
+
+def flash_attention_varlen(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    do_causal: bool = False,
+    scale: Optional[float] = None,
+    softcapping: float = 1.0,
+) -> None:
+    """
+    Flash Attention with variable-length sequences.
+    
+    Args:
+        out: Output tensor of shape [total_q_tokens, num_heads, head_dim]
+        query: Query tensor of shape [total_q_tokens, num_heads, head_dim]
+        key: Key tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        value: Value tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        cu_seqlens_q: Cumulative sequence lengths for queries, shape [batch_size + 1], dtype must be torch.int32
+        cu_seqlens_k: Cumulative sequence lengths for keys, shape [batch_size + 1], dtype must be torch.int32
+        max_seqlen_q: Maximum sequence length in the query batch
+        max_seqlen_k: Maximum sequence length in the key batch
+        do_causal: Whether to apply causal masking
+        scale: Attention scale factor (default: 1/sqrt(head_dim))
+        softcapping: Softcapping value (default: 1.0, must be 1.0 for this implementation)
+    
+    Note:
+        - cu_seqlens_q and cu_seqlens_k must have dtype torch.int32 for Metal compatibility
+        - Supported head dimensions: 32, 64, 72, 80, 96, 128
+        - Masks are not supported
+    """
+    if scale is None:
+        scale = query.shape[-1] ** -0.5
+    
+    ops.flash_attention_varlen(
+        out,
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        do_causal,
+        scale,
+        softcapping,
+    )
+
+def flash_attn_varlen_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: tuple = (-1, -1),
+    alibi_slopes: Optional[torch.Tensor] = None,
+    deterministic: bool = False,
+    return_attn_probs: bool = False,
+) -> torch.Tensor:
+    """
+    Flash Attention function with API compatible with the original Flash Attention.
+    
+    Note: This implementation does not support:
+    - dropout
+    - window attention
+    - alibi slopes
+    - returning attention probabilities
+    """
+    if dropout_p > 0:
+        raise NotImplementedError("Dropout is not supported in this implementation")
+    if window_size != (-1, -1):
+        raise NotImplementedError("Window attention is not supported")
+    if alibi_slopes is not None:
+        raise NotImplementedError("ALiBi is not supported")
+    if return_attn_probs:
+        raise NotImplementedError("Returning attention probabilities is not supported")
+    
+    # Create output tensor
+    out = torch.empty_like(q)
+    
+    # Call the kernel
+    flash_attention_varlen(
+        out=out,
+        query=q,
+        key=k,
+        value=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        do_causal=causal,
+        scale=softmax_scale,
+        softcapping=1.0,
+    )
+    
+    return out
+
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+]

build/torch29-metal-aarch64-darwin/_metal_flash_sdpa_metal_8924527.abi3.so

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

build/torch29-metal-aarch64-darwin/_ops.py

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

build/torch29-metal-aarch64-darwin/metadata.json

@@ -0,0 +1,5 @@
+{
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": []
+}

build/torch29-metal-aarch64-darwin/metal_flash_sdpa/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+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")))

flake.lock

@@ -0,0 +1,168 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1747046372,
+        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-compat_2": {
+      "locked": {
+        "lastModified": 1733328505,
+        "narHash": "sha256-NeCCThCEP3eCl2l/+27kNNK7QrwZB1IJCrXfrbv5oqU=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "ff81ac966bb2cae68946d5ed5fc4994f96d0ffec",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "flake-utils_2": {
+      "inputs": {
+        "systems": "systems_2"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "hf-nix": {
+      "inputs": {
+        "flake-compat": "flake-compat_2",
+        "flake-utils": "flake-utils_2",
+        "nixpkgs": "nixpkgs"
+      },
+      "locked": {
+        "lastModified": 1751968576,
+        "narHash": "sha256-cmKrlWpNTG/hq1bCaHXfbdm9T+Y6V+5//EHAVc1TLBE=",
+        "owner": "huggingface",
+        "repo": "hf-nix",
+        "rev": "3fcd1e1b46da91b6691261640ffd6b7123d0cb9e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "hf-nix",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "hf-nix": "hf-nix",
+        "nixpkgs": [
+          "kernel-builder",
+          "hf-nix",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1752758138,
+        "narHash": "sha256-twHfCnTbiVz1ScEOvtaJxjBJhqlx56YWhZOgg83qxzg=",
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "rev": "39cdfec039c0d2d57aeb4bfd7bb75fcf13187103",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1747820358,
+        "narHash": "sha256-fTqsZsUX6M3yeEvgyQvXcbGmT2CaRVyVwsi8eK29Oj4=",
+        "owner": "danieldk",
+        "repo": "nixpkgs",
+        "rev": "d3c1681180717528068082103bf323147de6ab0b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "danieldk",
+        "ref": "cudatoolkit-12.9-kernel-builder",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    },
+    "systems_2": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for SDPA kernel";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernel-builder";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genFlakeOutputs {
+      path = ./.;
+      rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;
+    };
+}

sdpa-metal/common.h

@@ -0,0 +1,7 @@
+#ifndef SDPA_METAL_COMMON_H
+#define SDPA_METAL_COMMON_H
+
+// Common definitions for Metal kernels
+// This file is included by Metal shaders, so it should not contain C++ code
+
+#endif // SDPA_METAL_COMMON_H

sdpa-metal/scaled_dot_product_attention.metal

@@ -0,0 +1,2101 @@
+#include <metal_stdlib>
+#include <metal_simdgroup>
+
+using namespace metal;
+
+#define STEEL_CONST static constant constexpr const
+#define STEEL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)")
+
+#if defined(__HAVE_BFLOAT__)
+
+typedef bfloat bfloat16_t;
+typedef half float16_t;
+
+#else
+
+typedef half float16_t;
+
+/////////////////////////////////////////////////////////////////////////////
+// Helpers
+/////////////////////////////////////////////////////////////////////////////
+
+constexpr METAL_FUNC uint16_t float_to_bfloat_bits(float x) {
+  // Check for nan
+  if ((as_type<uint32_t>(x) & ~_fp_encoding_traits<float>::sign_mask) >
+      _fp_encoding_traits<float>::inf_mask) {
+    return uint16_t(as_type<uint32_t>(0x7FC0));
+  }
+  // Take bits
+  uint32_t float_bits = as_type<uint32_t>(x);
+
+  // Round to nearest even
+  float_bits += ((float_bits >> 16) & 1) + as_type<uint32_t>(0x7FFF);
+
+  // Take upper 16 bits
+  return float_bits >> 16;
+}
+
+constexpr METAL_FUNC float bfloat_bits_to_float(uint16_t x) {
+  // Upper 16 bits are the data and lower 16 bits are 0s
+  return as_type<float>((uint32_t)x << 16);
+}
+
+struct _MLX_BFloat16;
+
+template <typename T>
+static constexpr constant bool can_convert_to_bfloat =
+    !is_same_v<T, _MLX_BFloat16> && is_convertible_v<T, float>;
+
+template <typename T>
+static constexpr constant bool can_convert_from_bfloat =
+    !is_same_v<T, _MLX_BFloat16> && is_convertible_v<float, T>;
+
+/////////////////////////////////////////////////////////////////////////////
+// Bfloat struct
+/////////////////////////////////////////////////////////////////////////////
+
+struct _MLX_BFloat16 {
+  /////////////////////////////////////////////////////////////////////////////
+  // Constructors
+  uint16_t bits_;
+  _MLX_BFloat16() thread = default;
+  _MLX_BFloat16() threadgroup = default;
+  _MLX_BFloat16() device = default;
+  _MLX_BFloat16() constant = default;
+
+  struct bits_to_bfloat_struct {};
+  static constexpr METAL_FUNC bits_to_bfloat_struct bits_to_bfloat() {
+    return bits_to_bfloat_struct();
+  }
+  constexpr METAL_FUNC _MLX_BFloat16(uint16_t bits, bits_to_bfloat_struct)
+      : bits_(bits) {}
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Conversions to bfloat
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_to_bfloat<T>>::type>
+  constexpr METAL_FUNC _MLX_BFloat16(T x) thread
+      : bits_(float_to_bfloat_bits(static_cast<float>(x))) {}
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_to_bfloat<T>>::type>
+  constexpr METAL_FUNC _MLX_BFloat16(T x) threadgroup
+      : bits_(float_to_bfloat_bits(static_cast<float>(x))) {}
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_to_bfloat<T>>::type>
+  constexpr METAL_FUNC _MLX_BFloat16(T x) device
+      : bits_(float_to_bfloat_bits(static_cast<float>(x))) {}
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_to_bfloat<T>>::type>
+  constexpr METAL_FUNC _MLX_BFloat16(T x) constant
+      : bits_(float_to_bfloat_bits(static_cast<float>(x))) {}
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Conversions from bfloat
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_from_bfloat<T>>::type>
+  constexpr METAL_FUNC operator T() const thread {
+    return static_cast<T>(bfloat_bits_to_float(bits_));
+  }
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_from_bfloat<T>>::type>
+  constexpr METAL_FUNC operator T() const threadgroup {
+    return static_cast<T>(bfloat_bits_to_float(bits_));
+  }
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_from_bfloat<T>>::type>
+  constexpr METAL_FUNC operator T() const device {
+    return static_cast<T>(bfloat_bits_to_float(bits_));
+  }
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_from_bfloat<T>>::type>
+  constexpr METAL_FUNC operator T() const constant {
+    return static_cast<T>(bfloat_bits_to_float(bits_));
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////
+// Bfloat operators
+/////////////////////////////////////////////////////////////////////////////
+
+/////////////////////////////////////////////////////////////////////////////
+// Unary ops
+constexpr METAL_FUNC _MLX_BFloat16 operator-(_MLX_BFloat16 x) {
+  return -static_cast<float>(x);
+}
+
+/////////////////////////////////////////////////////////////////////////////
+// Binary operators
+#define bfloat_binop_base(__op__, __operator__, otype, atype, btype, ctype) \
+  constexpr METAL_FUNC otype __operator__(atype lhs, btype rhs) {           \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);          \
+  }
+
+#define bfloat_binop_helper(__op__, __operator__, otype, itype, ctype)    \
+  constexpr METAL_FUNC otype __operator__(_MLX_BFloat16 lhs, itype rhs) { \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);        \
+  }                                                                       \
+  constexpr METAL_FUNC otype __operator__(itype lhs, _MLX_BFloat16 rhs) { \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);        \
+  }
+
+/////////////////////////////////////////////////////////////////////////////
+// Arithmetic Operators
+#define bfloat_binop(_op_, _operator_)                                       \
+  bfloat_binop_base(                                                         \
+      _op_, _operator_, _MLX_BFloat16, _MLX_BFloat16, _MLX_BFloat16, float); \
+  bfloat_binop_helper(_op_, _operator_, float, float, float);                \
+  bfloat_binop_helper(_op_, _operator_, float, half, float);                 \
+  bfloat_binop_helper(_op_, _operator_, _MLX_BFloat16, int32_t, float);      \
+  bfloat_binop_helper(_op_, _operator_, _MLX_BFloat16, uint32_t, float);     \
+  bfloat_binop_helper(_op_, _operator_, _MLX_BFloat16, int64_t, float);      \
+  bfloat_binop_helper(_op_, _operator_, _MLX_BFloat16, uint64_t, float);
+
+bfloat_binop(+, operator+);
+bfloat_binop(-, operator-);
+bfloat_binop(*, operator*);
+bfloat_binop(/, operator/);
+
+/////////////////////////////////////////////////////////////////////////////
+// Comparison ops
+#define bfloat_compop(__op__, __operator__)                             \
+  bfloat_binop_base(                                                    \
+      __op__, __operator__, bool, _MLX_BFloat16, _MLX_BFloat16, float); \
+  bfloat_binop_helper(__op__, __operator__, bool, float, float);        \
+  bfloat_binop_helper(__op__, __operator__, bool, half, float);         \
+  bfloat_binop_helper(__op__, __operator__, bool, int32_t, float);      \
+  bfloat_binop_helper(__op__, __operator__, bool, uint32_t, float);     \
+  bfloat_binop_helper(__op__, __operator__, bool, int64_t, float);      \
+  bfloat_binop_helper(__op__, __operator__, bool, uint64_t, float);
+
+bfloat_compop(>, operator>);
+bfloat_compop(<, operator<);
+bfloat_compop(>=, operator>=);
+bfloat_compop(<=, operator<=);
+bfloat_compop(==, operator==);
+bfloat_compop(!=, operator!=);
+
+#undef bfloat_compop
+#undef bfloat_binop_base
+#undef bfloat_binop_helper
+#undef bfloat_binop
+
+/////////////////////////////////////////////////////////////////////////////
+// Inplace Operators
+#define bfloat_inplace_op_helper(__op__, __operator__, itype, addr_space) \
+  constexpr METAL_FUNC addr_space _MLX_BFloat16& __operator__(            \
+      addr_space _MLX_BFloat16& lhs, itype rhs) {                         \
+    lhs = static_cast<float>(lhs) __op__ static_cast<float>(rhs);         \
+    return lhs;                                                           \
+  }                                                                       \
+  constexpr METAL_FUNC addr_space itype& __operator__(                    \
+      addr_space itype& lhs, _MLX_BFloat16 rhs) {                         \
+    lhs = static_cast<float>(lhs) __op__ static_cast<float>(rhs);         \
+    return lhs;                                                           \
+  }
+
+#define bfloat_inplace_op_addr_space_helper(__op__, __operator__, itype) \
+  bfloat_inplace_op_helper(__op__, __operator__, itype, device);         \
+  bfloat_inplace_op_helper(__op__, __operator__, itype, thread);         \
+  bfloat_inplace_op_helper(__op__, __operator__, itype, threadgroup);
+
+#define bfloat_inplace_op(itype)                             \
+  bfloat_inplace_op_addr_space_helper(+, operator+=, itype); \
+  bfloat_inplace_op_addr_space_helper(-, operator-=, itype); \
+  bfloat_inplace_op_addr_space_helper(*, operator*=, itype); \
+  bfloat_inplace_op_addr_space_helper(/, operator/=, itype);
+
+bfloat_inplace_op(float);
+bfloat_inplace_op(half);
+bfloat_inplace_op(int16_t);
+bfloat_inplace_op(int32_t);
+bfloat_inplace_op(int64_t);
+bfloat_inplace_op(uint16_t);
+bfloat_inplace_op(uint32_t);
+bfloat_inplace_op(uint64_t);
+
+#undef bfloat_inplace_op_helper
+#undef bfloat_inplace_op_addr_space_helper
+#undef bfloat_inplace_op
+
+#define bfloat_inplace_op_helper(__op__, __operator__, addr_space) \
+  constexpr METAL_FUNC addr_space _MLX_BFloat16& __operator__(     \
+      addr_space _MLX_BFloat16& lhs, _MLX_BFloat16 rhs) {          \
+    lhs = static_cast<float>(lhs) __op__ static_cast<float>(rhs);  \
+    return lhs;                                                    \
+  }
+
+#define bfloat_inplace_op_addr_space_helper(__op__, __operator__) \
+  bfloat_inplace_op_helper(__op__, __operator__, device);         \
+  bfloat_inplace_op_helper(__op__, __operator__, thread);         \
+  bfloat_inplace_op_helper(__op__, __operator__, threadgroup);
+
+bfloat_inplace_op_addr_space_helper(+, operator+=);
+bfloat_inplace_op_addr_space_helper(-, operator-=);
+bfloat_inplace_op_addr_space_helper(*, operator*=);
+bfloat_inplace_op_addr_space_helper(/, operator/=);
+
+#undef bfloat_inplace_op_helper
+#undef bfloat_inplace_op_addr_space_helper
+
+/////////////////////////////////////////////////////////////////////////////
+// Bfloat typedef
+/////////////////////////////////////////////////////////////////////////////
+
+typedef struct _MLX_BFloat16 bfloat16_t;
+
+#endif
+
+// ============ "mlx/backend/metal/kernels/scaled_dot_product_attention_params.h"
+
+struct MLXFastAttentionParams {
+  const int M;
+  const int N;
+  const int K;
+
+  const int ldq; // ldq == ldo
+  const int ldk;
+  const int ldv;
+  const int lds;
+  const int ldo;
+
+  const int tiles_n;
+  const int tiles_m;
+
+  const int batch_stride_q;
+  const int batch_stride_k;
+  const int batch_stride_v;
+  const int batch_stride_o;
+
+  const int swizzle_log;
+  const int gemm_n_iterations_aligned;
+  const int gemm_k_iterations_aligned;
+  const int gemm_sv_m_block_iterations;
+
+  const int batch_ndim;
+  const float alpha;
+  const float softcapping;
+};
+
+struct MLXScaledDotProductAttentionParams {
+  // Associated dimensions & transposition information
+  const uint QUERY_SEQUENCE_LENGTH = 1;
+  const uint N_Q_HEADS = 32;
+  const uint N_KV_HEADS = 32;
+  const uint KV_TILES = 1;
+  const float INV_ALPHA = 0.08838834764831843f;
+};
+
+// ============ "mlx/backend/metal/kernels/scaled_dot_product_attention_params.sdpa_vector"
+
+
+// ============ "mlx/backend/metal/kernels/utils.h"
+
+template <typename U>
+struct Limits {
+  static const constant U max = metal::numeric_limits<U>::max();
+  static const constant U min = metal::numeric_limits<U>::min();
+  static const constant U finite_max = metal::numeric_limits<U>::max();
+  static const constant U finite_min = metal::numeric_limits<U>::min();
+};
+
+#define instantiate_default_limit(type)                                      \
+  template <>                                                                \
+  struct Limits<type> {                                                      \
+    static constexpr constant type max = metal::numeric_limits<type>::max(); \
+    static constexpr constant type min = metal::numeric_limits<type>::min(); \
+    static constexpr constant type finite_max =                              \
+        metal::numeric_limits<type>::max();                                  \
+    static constexpr constant type finite_min =                              \
+        metal::numeric_limits<type>::min();                                  \
+  };
+
+instantiate_default_limit(uint8_t);
+instantiate_default_limit(uint16_t);
+instantiate_default_limit(uint32_t);
+instantiate_default_limit(uint64_t);
+instantiate_default_limit(int8_t);
+instantiate_default_limit(int16_t);
+instantiate_default_limit(int32_t);
+instantiate_default_limit(int64_t);
+
+#define instantiate_float_limit(type)             \
+  template <>                                     \
+  struct Limits<type> {                           \
+    static constexpr constant type max =          \
+        metal::numeric_limits<type>::infinity();  \
+    static constexpr constant type min =          \
+        -metal::numeric_limits<type>::infinity(); \
+    static constexpr constant type finite_max =   \
+        metal::numeric_limits<type>::max();       \
+    static constexpr constant type finite_min =   \
+        -metal::numeric_limits<type>::max();      \
+  };
+
+instantiate_float_limit(half);
+instantiate_float_limit(float);
+instantiate_float_limit(bfloat16_t);
+
+
+// ============ "mlx/backend/metal/kernels/steel/attn/loader.h"
+
+template <
+    typename T,
+    short BROWS,
+    short BCOLS,
+    short dst_ld,
+    short reduction_dim,
+    short tgp_size,
+    short alignment = 1,
+    short n_reads = (BCOLS * BROWS) / (tgp_size),
+    short TCOLS = BCOLS / n_reads,
+    short TROWS = tgp_size / TCOLS>
+struct BlockLoader {
+  STEEL_CONST short n_rows = (BROWS + TROWS - 1) / TROWS;
+  STEEL_CONST short vec_size = n_reads;
+
+  // Leading dimension for src
+  const int src_ld;
+  const int tile_stride;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  struct alignas(alignment * sizeof(T)) ReadVector {
+    uint8_t v[sizeof(T) * vec_size];
+  };
+
+  /* Constructor */
+  METAL_FUNC BlockLoader(
+      const device T* src_,
+      const int src_ld_,
+      threadgroup T* dst_,
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(src_ld_),
+        tile_stride(reduction_dim ? BCOLS : BROWS * src_ld),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        src(src_ + bi * src_ld + bj) {}
+
+  /* Apply operation to threadgroup without bound checking */
+  template <typename UnaryOp>
+  METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * dst_ld + j] = op.apply(dst[i * dst_ld + j]);
+      }
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      *((threadgroup ReadVector*)(&dst[i * dst_ld])) =
+          *((const device ReadVector*)(&src[i * src_ld]));
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - with bound checking */
+  METAL_FUNC void load_safe(short2 src_tile_dim) const {
+    src_tile_dim = src_tile_dim - short2(bj, bi);
+
+    // Skip loading if thread has no valid reads
+    if (src_tile_dim.x <= 0 || src_tile_dim.y <= 0) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = T(0);
+        }
+      }
+      return;
+    }
+
+    // Use fast thread memory for bound checks
+    bool tmp_idx[vec_size];
+    T tmp_val[vec_size];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      // Make sure tmp_idx only contains valid indices
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_idx[j] = (i < src_tile_dim.y) && (j < src_tile_dim.x);
+      }
+
+      // Read valid indices into tmp_val
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = src[(tmp_idx[j] ? i * src_ld + j : 0)];
+      }
+
+      // Zero out uneeded values
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = tmp_idx[j] ? tmp_val[j] : T(0);
+      }
+
+      // Copy values to threadgroup memory
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * dst_ld + j] = tmp_val[j];
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    src += tile_stride;
+  }
+};
+
+template <int R, int C>
+struct CShape {
+  STEEL_CONST int kRows = R;
+  STEEL_CONST int kCols = C;
+};
+
+template <
+    typename T,
+    short BROWS,
+    short BCOLS,
+    short kDstStrRow,
+    short kDstStrCol,
+    short reduction_dim,
+    short tgp_size,
+    short n_reads = (BCOLS * BROWS) / (tgp_size),
+    short TCOLS = BCOLS / n_reads,
+    short TROWS = tgp_size / TCOLS>
+struct BlockLoaderT {
+  STEEL_CONST short n_rows = (BROWS + TROWS - 1) / TROWS;
+  STEEL_CONST short vec_size = n_reads;
+
+  // Leading dimension for src
+  const int src_ld;
+  const int tile_stride;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  /* Constructor */
+  METAL_FUNC BlockLoaderT(
+      const device T* src_,
+      const int src_ld_,
+      threadgroup T* dst_,
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(src_ld_),
+        tile_stride(reduction_dim ? BCOLS : BROWS * src_ld),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * kDstStrRow + bj * kDstStrCol),
+        src(src_ + bi * src_ld + bj) {}
+
+  /* Apply operation to threadgroup without bound checking */
+  template <typename UnaryOp>
+  METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * kDstStrRow + j * kDstStrCol] =
+            op.apply(dst[i * kDstStrRow + j * kDstStrCol]);
+      }
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * kDstStrRow + j * kDstStrCol] = src[i * src_ld + j];
+      }
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - with bound checking */
+  METAL_FUNC void load_safe(short2 src_tile_dim) const {
+    src_tile_dim = src_tile_dim - short2(bj, bi);
+
+    // Skip loading if thread has no valid reads
+    if (src_tile_dim.x <= 0 || src_tile_dim.y <= 0) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * kDstStrRow + j * kDstStrCol] = T(0);
+        }
+      }
+      return;
+    }
+
+    // Use fast thread memory for bound checks
+    bool tmp_idx[vec_size];
+    T tmp_val[vec_size];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      // Make sure tmp_idx only contains valid indices
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_idx[j] = (i < src_tile_dim.y) && (j < src_tile_dim.x);
+      }
+
+      // Read valid indices into tmp_val
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = src[(tmp_idx[j] ? i * src_ld + j : 0)];
+      }
+
+      // Zero out uneeded values
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = tmp_idx[j] ? tmp_val[j] : T(0);
+      }
+
+      // Copy values to threadgroup memory
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * kDstStrRow + j * kDstStrCol] = tmp_val[j];
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    src += tile_stride;
+  }
+};
+
+// ============ "mlx/backend/metal/kernels/steel/utils/type_traits.h"
+
+template <typename... Ts>
+struct make_void {
+  typedef void type;
+};
+
+template <typename... Ts>
+using void_t = typename make_void<Ts...>::type;
+
+template <typename T>
+struct pointer_element {};
+
+template <typename T>
+struct pointer_element<thread T*> {
+  using type = remove_cv_t<T>;
+};
+template <typename T>
+struct pointer_element<device T*> {
+  using type = remove_cv_t<T>;
+};
+template <typename T>
+struct pointer_element<constant T*> {
+  using type = remove_cv_t<T>;
+};
+template <typename T>
+struct pointer_element<threadgroup T*> {
+  using type = remove_cv_t<T>;
+};
+
+template <typename T>
+using pointer_element_t = typename pointer_element<remove_cv_t<T>>::type;
+
+// ============ "mlx/backend/metal/kernels/steel/utils/integral_constant.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Integral constant with casting
+///////////////////////////////////////////////////////////////////////////////
+
+template <int val>
+using Int = integral_constant<int, val>;
+
+///////////////////////////////////////////////////////////////////////////////
+// Binary Operators on Integral constants
+///////////////////////////////////////////////////////////////////////////////
+
+#define integral_const_binop(__op__, __operator__)          \
+  template <typename T, T tv, typename U, U uv>             \
+  METAL_FUNC constexpr auto __operator__(                   \
+      integral_constant<T, tv>, integral_constant<U, uv>) { \
+    constexpr auto res = tv __op__ uv;                      \
+    return integral_constant<decltype(res), res>{};         \
+  }
+
+integral_const_binop(+, operator+);
+integral_const_binop(-, operator-);
+integral_const_binop(*, operator*);
+integral_const_binop(/, operator/);
+
+integral_const_binop(==, operator==);
+integral_const_binop(!=, operator!=);
+integral_const_binop(<, operator<);
+integral_const_binop(>, operator>);
+integral_const_binop(<=, operator<=);
+integral_const_binop(>=, operator>=);
+
+integral_const_binop(&&, operator&&);
+integral_const_binop(||, operator||);
+
+#undef integral_const_binop
+
+///////////////////////////////////////////////////////////////////////////////
+// Reduction operators
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+METAL_FUNC constexpr T sum(T x) {
+  return x;
+}
+
+template <typename T, typename... Us>
+METAL_FUNC constexpr auto sum(T x, Us... us) {
+  return x + sum(us...);
+}
+
+// ============ "mlx/backend/metal/kernels/steel/gemm/transforms.h"
+
+template <typename OutT, typename InT>
+struct TransformNone {
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  static METAL_FUNC OutT apply(InT x, OutT) {
+    return static_cast<OutT>(x);
+  }
+};
+
+template <typename OutT, typename InT>
+struct TransformAdd {
+  TransformAdd(const float, const float) {}
+
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  static METAL_FUNC OutT apply(InT x, OutT c) {
+    return static_cast<OutT>(x) + c;
+  }
+};
+
+template <typename OutT, typename InT>
+struct TransformAxpby {
+  const float alpha;
+  const float beta;
+
+  TransformAxpby(const float alpha_, const float beta_)
+      : alpha(alpha_), beta(beta_) {}
+
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  METAL_FUNC OutT apply(InT x, OutT c) const {
+    return static_cast<OutT>(x * alpha + (beta * c));
+  }
+};
+
+template <typename T>
+struct AccumHelper {
+  typedef float accum_type;
+};
+
+struct BlockSwizzle {
+  static METAL_FUNC int2
+  swizzle(uint3 tid [[threadgroup_position_in_grid]], const int swizzle_log) {
+    const int tid_x = (tid.x) >> swizzle_log;
+    const int tid_y =
+        ((tid.y) << swizzle_log) + ((tid.x) & ((1 << swizzle_log) - 1));
+    return int2(tid_x, tid_y);
+  }
+};
+
+// ============ "mlx/backend/metal/kernels/steel/attn/mma.h"
+
+template <typename RInt, typename CInt>
+struct Shape2D {
+  RInt r;
+  CInt c;
+
+  Shape2D(RInt r_, CInt c_) : r(r_), c(c_) {}
+};
+
+template <typename Shape, typename Layout>
+struct Layout2D {
+  Shape shape;
+  Layout layout;
+};
+
+template <typename T, int kFragRows_, int kFragCols_>
+struct BaseMMAFrag {
+  static_assert(
+      kFragRows_ == 8,
+      "Only 8 x 8 fragment matrices are currently supported");
+  static_assert(
+      kFragCols_ == 8,
+      "Only 8 x 8 fragment matrices are currently supported");
+};
+
+template <typename T>
+struct BaseMMAFrag<T, 8, 8> {
+  STEEL_CONST int kFragRows = 8;
+  STEEL_CONST int kFragCols = 8;
+
+  STEEL_CONST int kElemsPerFrag = (kFragRows * kFragCols) / 32;
+
+  STEEL_CONST int kElemRows = 1;
+  STEEL_CONST int kElemCols = 2;
+
+  static_assert(
+      kElemRows * kElemCols == kElemsPerFrag,
+      "MMAFrag shape is not consistent with MMAFrag size");
+
+  typedef metal::simdgroup_matrix<T, kFragRows, kFragCols> mat_type;
+  typedef metal::vec<T, kElemsPerFrag> frag_type;
+  typedef metal::vec<T, kElemRows> row_frag_type;
+  typedef metal::vec<T, kElemCols> col_frag_type;
+
+  template <typename U>
+  using dtype_mat_t = typename metal::simdgroup_matrix<U, kFragRows, kFragCols>;
+
+  template <typename U>
+  using dtype_frag_t = typename metal::vec<U, kElemsPerFrag>;
+
+  METAL_FUNC static constexpr short2 get_coord(ushort simd_lane_id
+                                               [[thread_index_in_simdgroup]]) {
+    const short qid = simd_lane_id / 4;
+    const short fm = (qid & 4) + ((simd_lane_id / 2) % 4);
+    const short fn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
+    return short2{fn, fm};
+  }
+
+  template <typename SrcPtrType, typename StrX, typename StrY>
+  METAL_FUNC static constexpr void
+  load(thread frag_type& dst, SrcPtrType src, StrX str_x, StrY str_y) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        dst[i * kElemCols + j] = static_cast<T>(src[i * str_x.value + j * str_y.value]);
+      }
+    }
+  }
+
+  template <
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX,
+      typename OffY>
+  METAL_FUNC static constexpr void load_safe(
+      thread frag_type& dst,
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = Int<0>{},
+      OffY off_y = Int<0>{}) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        if ((off_x + i) < lim_x && (off_y + j) < lim_y) {
+          dst[i * kElemCols + j] =
+              static_cast<T>(src[(off_x + i) * str_x + (off_y + j) * str_y.value]);
+        } else {
+          dst[i * kElemCols + j] = T(0);
+        }
+      }
+    }
+  }
+
+  template <typename DstPtrType, typename StrX, typename StrY>
+  METAL_FUNC static constexpr void
+  store(const thread frag_type& src, DstPtrType dst, StrX str_x, StrY str_y) {
+    using U = pointer_element_t<DstPtrType>;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        dst[i * str_x + j * str_y.value] = static_cast<U>(src[i * kElemCols + j]);
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX,
+      typename OffY>
+  METAL_FUNC static constexpr void store_safe(
+      const thread frag_type& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = Int<0>{},
+      OffY off_y = Int<0>{}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        if ((off_x + i) < lim_x && (off_y + j) < lim_y) {
+          dst[(off_x + i) * str_x + (off_y + j) * str_y.value] =
+              static_cast<U>(src[i * kElemCols + j]);
+        }
+      }
+    }
+  }
+
+  template <typename Atype, typename Btype, typename Ctype>
+  METAL_FUNC static constexpr void mma(
+      thread frag_type& D,
+      thread dtype_frag_t<Atype>& A,
+      thread dtype_frag_t<Btype>& B,
+      thread dtype_frag_t<Ctype>& C) {
+    mat_type D_mat;
+    dtype_mat_t<Atype> A_mat;
+    dtype_mat_t<Btype> B_mat;
+    dtype_mat_t<Ctype> C_mat;
+
+    reinterpret_cast<thread dtype_frag_t<Atype>&>(A_mat.thread_elements()) = A;
+    reinterpret_cast<thread dtype_frag_t<Btype>&>(B_mat.thread_elements()) = B;
+    reinterpret_cast<thread dtype_frag_t<Ctype>&>(C_mat.thread_elements()) = C;
+
+    mma(D_mat, A_mat, B_mat, C_mat);
+
+    D = reinterpret_cast<thread frag_type&>(D_mat.thread_elements());
+  }
+
+  template <typename Atype, typename Btype, typename Ctype>
+  METAL_FUNC static constexpr void mma(
+      thread mat_type& D,
+      thread dtype_mat_t<Atype>& A,
+      thread dtype_mat_t<Btype>& B,
+      thread dtype_mat_t<Ctype>& C) {
+    simdgroup_multiply_accumulate(D, A, B, C);
+  }
+
+  template <typename Op>
+  METAL_FUNC static constexpr void row_reduce(
+      thread const frag_type& inp_vals,
+      thread T* reduced_vals) {
+    T thr_reduce = Op::apply(inp_vals.x, inp_vals.y);
+
+    T qgr_reduce = simd_shuffle_xor(thr_reduce, ushort(1));
+    qgr_reduce = Op::apply(thr_reduce, qgr_reduce);
+
+    T sgr_reduce = simd_shuffle_xor(qgr_reduce, ushort(8));
+    sgr_reduce = Op::apply(qgr_reduce, sgr_reduce);
+
+    reduced_vals[0] = Op::apply(reduced_vals[0], sgr_reduce);
+  }
+
+  template <typename Op>
+  METAL_FUNC static constexpr void row_bin_op(
+      thread frag_type& inp_vals,
+      thread T* row_vals) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        inp_vals[i * kElemCols + j] =
+            Op::apply(inp_vals[i * kElemCols + j], row_vals[i]);
+      }
+    }
+  }
+};
+
+template <
+    typename T,
+    int kTileRows_,
+    int kTileCols_,
+    class MMAFrag_ = BaseMMAFrag<T, 8, 8>>
+struct MMATile {
+  using MMAFrag_t = MMAFrag_;
+  using elem_type = T;
+  STEEL_CONST int kFragRows = MMAFrag_t::kFragRows;
+  STEEL_CONST int kFragCols = MMAFrag_t::kFragCols;
+  STEEL_CONST int kElemsPerFrag = MMAFrag_t::kElemsPerFrag;
+
+  STEEL_CONST int kTileRows = kTileRows_;
+  STEEL_CONST int kTileCols = kTileCols_;
+
+  STEEL_CONST int kRows = kTileRows * kFragRows;
+  STEEL_CONST int kCols = kTileCols * kFragCols;
+
+  STEEL_CONST int kNumFrags = kTileRows * kTileCols;
+  STEEL_CONST int kElemsPerTile = kNumFrags * kElemsPerFrag;
+
+  STEEL_CONST int kRowsPerThread = kTileRows * MMAFrag_t::kElemRows;
+  STEEL_CONST int kColsPerThread = kTileCols * MMAFrag_t::kElemCols;
+
+  typedef typename MMAFrag_t::mat_type mat_type;
+  typedef typename MMAFrag_t::frag_type frag_type;
+
+  frag_type val_frags[kNumFrags]; // = {frag_type(0)};
+
+  METAL_FUNC MMATile() thread {}
+
+  METAL_FUNC constexpr void clear() {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kNumFrags; ++i) {
+      val_frags[i] = frag_type(0);
+    }
+  }
+
+  METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) {
+    return val_frags[i * kTileCols + j];
+  }
+
+  METAL_FUNC constexpr const thread frag_type& frag_at(
+      const short i,
+      const short j) const {
+    return val_frags[i * kTileCols + j];
+  }
+
+  METAL_FUNC mat_type mat_at(const short i, const short j) {
+    mat_type val_mat;
+    STEEL_PRAGMA_UNROLL
+    for (short ii = 0; ii < kElemsPerFrag; ++ii) {
+      val_mat.thread_elements()[ii] = frag_at(i, j)[ii];
+    }
+    return val_mat;
+  }
+
+  METAL_FUNC thread elem_type* elems() {
+    return reinterpret_cast<thread elem_type*>(val_frags);
+  }
+
+  METAL_FUNC const thread elem_type* elems() const {
+    return reinterpret_cast<const thread elem_type*>(val_frags);
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_reduce(thread T vals[kRowsPerThread]) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::template row_reduce<Op>(
+            frag_at(i, j), &vals[i * MMAFrag_t::kElemRows]);
+      }
+    }
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_bin_op(thread T vals[kRowsPerThread]) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::template row_bin_op<Op>(
+            frag_at(i, j), &vals[i * MMAFrag_t::kElemRows]);
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y, int str_x, int str_y>
+  METAL_FUNC void load(const threadgroup U* src) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::load(
+            frag_at(i, j),
+            &(
+                src[(i * kFragRows) * w_x * str_x +
+                    (j * kFragCols) * w_y * str_y]),
+            Int<str_x>{},
+            Int<str_y>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y, int str_x, int str_y>
+  METAL_FUNC void store(threadgroup U* dst) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::store(
+            frag_at(i, j),
+            &(
+                dst[(i * kFragRows) * w_x * str_x +
+                    (j * kFragCols) * w_y * str_y]),
+            Int<str_x>{},
+            Int<str_y>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void load(const device U* src, const int ld) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::load(
+            frag_at(i, j),
+            &(src[(i * kFragRows) * w_x * ld + (j * kFragCols) * w_y]),
+            ld,
+            Int<1>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void store(device U* dst, const int ld) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::store(
+            frag_at(i, j),
+            &(dst[(i * kFragRows) * w_x * ld + (j * kFragCols) * w_y]),
+            ld,
+            Int<1>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void
+  load_safe(const device U* src, const int ld, const short2 src_tile_dims) {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::load_safe(
+            frag_at(i, j),
+            src,
+            ld,
+            Int<1>{},
+            src_tile_dims.y,
+            src_tile_dims.x,
+            (i * kFragRows) * w_x,
+            (j * kFragCols) * w_y);
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void
+  store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::store_safe(
+            frag_at(i, j),
+            dst,
+            ld,
+            Int<1>{},
+            dst_tile_dims.y,
+            dst_tile_dims.x,
+            (i * kFragRows) * w_x,
+            (j * kFragCols) * w_y);
+      }
+    }
+  }
+};
+
+template <
+    typename Dtype,
+    typename Atype,
+    typename Btype,
+    typename Ctype,
+    int M,
+    int N,
+    int K,
+    class MMAFragD,
+    class MMAFragA,
+    class MMAFragB,
+    class MMAFragC>
+METAL_FUNC void tile_matmad(
+    thread MMATile<Dtype, M, N, MMAFragD>& D,
+    thread MMATile<Atype, M, K, MMAFragA>& A,
+    thread MMATile<Btype, K, N, MMAFragB>& B,
+    thread MMATile<Ctype, M, N, MMAFragC>& C) {
+  STEEL_PRAGMA_UNROLL
+  for (short m = 0; m < M; ++m) {
+    STEEL_PRAGMA_UNROLL
+    for (short n = 0; n < N; ++n) {
+      short m_serp = m; //(n % 2) ? (M - 1 - m) : m;
+      short n_serp = (m % 2) ? (N - 1 - n) : n;
+
+      STEEL_PRAGMA_UNROLL
+      for (short k = 0; k < K; ++k) {
+        MMAFragD::mma(
+            D.frag_at(m_serp, n_serp),
+            A.frag_at(m_serp, k),
+            B.frag_at(k, n_serp),
+            C.frag_at(m_serp, n_serp));
+      }
+    }
+  }
+}
+
+template <
+    typename T,
+    typename U,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    short lda_tgp,
+    short ldb_tgp,
+    typename AccumType = float,
+    typename Epilogue = TransformNone<U, AccumType>>
+struct BlockMMA {
+  // MMAFrag size
+  STEEL_CONST short kFragSize = 8;
+  using MMAFrag_acc_t = BaseMMAFrag<AccumType, kFragSize, kFragSize>;
+
+  // Warp tile simdgroup matrix strides along M
+  STEEL_CONST short TM_stride = kFragSize * WM;
+  // Warp tile simdgroup matrix strides along M
+  STEEL_CONST short TN_stride = kFragSize * WN;
+
+  // Warp tile size along M
+  STEEL_CONST short TM = BM / TM_stride;
+  // Warp tile size along N
+  STEEL_CONST short TN = BN / TN_stride;
+
+  // Threadgroup A strides
+  STEEL_CONST short A_str_m = transpose_a ? 1 : lda_tgp; // M
+  STEEL_CONST short A_str_k = transpose_a ? lda_tgp : 1; // K
+
+  // Threadgroup B strides
+  STEEL_CONST short B_str_k = transpose_b ? 1 : ldb_tgp; // K
+  STEEL_CONST short B_str_n = transpose_b ? ldb_tgp : 1; // N
+
+  // Threadgroup strides along K
+  STEEL_CONST short tile_stride_a = kFragSize * A_str_k;
+  STEEL_CONST short tile_stride_b = kFragSize * B_str_k;
+
+  // Simdgroup matrices
+  MMATile<AccumType, TM, 1, MMAFrag_acc_t> Atile;
+  MMATile<AccumType, 1, TN, MMAFrag_acc_t> Btile;
+  MMATile<AccumType, TM, TN, MMAFrag_acc_t> Ctile;
+
+  // Offsets within threadgroup
+  short sm;
+  short sn;
+
+  short As_offset;
+  short Bs_offset;
+
+  /* Constructor */
+  METAL_FUNC BlockMMA(
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]]) {
+    // Determine thread position in simdgroup matrix
+    short tm = kFragSize * (simd_group_id / WN);
+    short tn = kFragSize * (simd_group_id % WN);
+
+    short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id);
+    sm = simd_coord.y;
+    sn = simd_coord.x;
+
+    // Determine thread and simdgroup offset
+    As_offset = (tm + sm) * A_str_m + (sn)*A_str_k; // M, K
+    Bs_offset = (sm)*B_str_k + (tn + sn) * B_str_n; // K, N
+
+    sm += tm;
+    sn += tn;
+  }
+
+  /* (BM, BK) X (BK, BN) multiply accumulate function */
+  METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) {
+    // Adjust for simdgroup and thread location
+    As += As_offset;
+    Bs += Bs_offset;
+
+    // Iterate over BK in blocks of kFragSize
+    STEEL_PRAGMA_UNROLL
+    for (short kk = 0; kk < BK; kk += kFragSize) {
+      simdgroup_barrier(mem_flags::mem_none);
+
+      Atile.template load<T, WM, 1, A_str_m, A_str_k>(As);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      Btile.template load<T, 1, WN, B_str_k, B_str_n>(Bs);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      tile_matmad(Ctile, Atile, Btile, Ctile);
+
+      // Progress to next simdgroup tile
+      As += tile_stride_a;
+      Bs += tile_stride_b;
+    }
+  }
+
+  /* Store results from simdgroup_matrix results into device memory */
+  METAL_FUNC void store_result(device U* D, const int ldd) {
+    // Apply epilogue
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
+      Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]);
+    }
+
+    // Adjust for simdgroup and thread location
+    D += sm * ldd + sn;
+
+    Ctile.template store<U, WM, WN>(D, ldd);
+  }
+
+  METAL_FUNC void
+  store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) {
+    // Apply epilogue
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
+      Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]);
+    }
+
+    // Adjust for simdgroup and thread location
+    D += sm * ldd + sn;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    Ctile.template store_safe<U, WM, WN>(D, ldd, dst_tile_dims);
+  }
+
+  /* Apply epilogue */
+  template <typename UnaryEpilogue>
+  METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) {
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
+      Ctile.elems()[i] = epilogue_op.apply(Ctile.elems()[i]);
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename BinaryEpilogue>
+  METAL_FUNC void apply_epilogue(
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      thread const BinaryEpilogue& epilogue_op) {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in C
+        thread auto& accum = Ctile.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < decltype(Ctile)::kElemsPerFrag; k++) {
+          accum[k] = epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
+        }
+      }
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename BinaryEpilogue>
+  METAL_FUNC void apply_epilogue_safe(
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      short2 dst_tile_dims,
+      thread const BinaryEpilogue& epilogue_op) {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in C
+        thread auto& accum = Ctile.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+
+        constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
+
+        // Read C
+        U c_elems[kelems] = {0};
+
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          if ((j * TN_stride + k) < dst_tile_dims.x) {
+            c_elems[k] = C[offset_c + k * fdc];
+          }
+        }
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          accum[k] = epilogue_op.apply(accum[k], c_elems[k]);
+        }
+      }
+    }
+  }
+
+  /* Store results from simdgroup_matrix results into device memory */
+  METAL_FUNC void store_result(
+      device U* D,
+      const int ldd,
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      thread const Epilogue& epilogue_op) const {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    D += (sm)*ldd + sn;
+
+    constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in C
+        thread const auto& accum = Ctile.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+        int offset_d = (i * TM_stride) * ldd + (j * TN_stride);
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          D[offset_d + k] = epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
+        }
+      }
+    }
+  }
+
+  METAL_FUNC void store_result_safe(
+      device U* D,
+      const int ldd,
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      short2 dst_tile_dims,
+      thread const Epilogue& epilogue_op) const {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    D += (sm)*ldd + sn;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
+
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < TM; i++) {
+      if (i * TM_stride < dst_tile_dims.y) {
+        STEEL_PRAGMA_UNROLL
+        for (int j = 0; j < TN; j++) {
+          // Get accumulated result and associated offset in C
+          thread const auto& accum = Ctile.frag_at(i, j);
+          int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+          int offset_d = (i * TM_stride) * ldd + (j * TN_stride);
+
+          // Apply epilogue
+          STEEL_PRAGMA_UNROLL
+          for (short k = 0; k < kelems; k++) {
+            if ((j * TN_stride + k) < dst_tile_dims.x) {
+              D[offset_d + k] =
+                  epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
+            }
+          }
+        }
+      }
+    }
+  }
+};
+
+// ============ "mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h"
+
+struct AttnParams {
+  int B; ///< Batch Size
+  int H; ///< Heads
+  int D; ///< Head Dim
+
+  int qL; ///< Query Sequence Length
+  int kL; ///< Key Sequence Length
+
+  int gqa_factor; ///< Group Query factor
+  float scale; ///< Attention scale
+  float softcapping; ///< Softcapping value (1.0 for no softcapping)
+
+  int NQ; ///< Number of query blocks
+  int NK; ///< Number of key/value blocks
+
+  int NQ_aligned; ///< Number of full query blocks
+  int NK_aligned; ///< Number of full key/value blocks
+
+  int qL_rem; ///< Remainder in last query block
+  int kL_rem; ///< Remainder in last key/value block
+  int qL_off; ///< Offset in query sequence start
+
+  int64_t Q_strides[3]; ///< Query  strides (B, H, L, D = 1)
+  int64_t K_strides[3]; ///< Key    strides (B, H, L, D = 1)
+  int64_t V_strides[3]; ///< Value  strides (B, H, L, D = 1)
+  int64_t O_strides[3]; ///< Output strides (B, H, L, D = 1)
+  
+  // Flash Attention variable-length support
+  int total_q_tokens; ///< Total number of query tokens (sum of all sequence lengths)
+  int total_k_tokens; ///< Total number of key/value tokens
+  int max_seqlen_q; ///< Maximum query sequence length
+  int max_seqlen_k; ///< Maximum key/value sequence length
+};
+
+struct AttnMaskParams {
+  int64_t M_strides[3]; ///< Mask  strides (B, H, qL, kL = 1)
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM kernels
+///////////////////////////////////////////////////////////////////////////////
+
+constant bool align_Q [[function_constant(200)]];
+constant bool align_K [[function_constant(201)]];
+
+constant bool has_mask [[function_constant(300)]];
+constant bool do_causal [[function_constant(301)]];
+
+template <typename T>
+struct TransformScale {
+  T scale;
+  METAL_FUNC TransformScale(T scale_) : scale(scale_) {}
+
+  METAL_FUNC T apply(T x) const {
+    return scale * x;
+  }
+};
+
+struct MaxOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return metal::max(x, y);
+  }
+};
+
+struct SumOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x + y;
+  }
+};
+
+struct MulOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x * y;
+  }
+};
+
+struct SubOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x - y;
+  }
+};
+
+struct ExpSubOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return fast::exp2(x - y);
+  }
+};
+
+struct DivOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x / y;
+  }
+};
+
+// clang-format off
+template <
+    typename T,
+    int BQ,
+    int BK,
+    int BD,
+    int WM,
+    int WN,
+    typename MaskType = float,
+    typename AccumType = float>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void attention(
+    const device T* Q [[buffer(0)]],
+    const device T* K [[buffer(1)]],
+    const device T* V [[buffer(2)]],
+    device T* O [[buffer(3)]],
+    const constant AttnParams* params [[buffer(4)]],
+    const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]],
+    const device MaskType* mask [[buffer(6), function_constant(has_mask)]],
+    const device int* cu_seqlens_q [[buffer(7)]],  // Cumulative query sequence lengths
+    const device int* cu_seqlens_k [[buffer(8)]],  // Cumulative key sequence lengths
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) { // clang-format on
+
+  // Pacifying compiler
+  (void)lid;
+
+  // Flash Attention variable-length indexing
+  // tid.z is now the sequence index within the batch
+  int batch_idx = tid.z;
+  int head_idx = tid.y;
+  int block_idx = tid.x;
+  
+  // Get sequence boundaries from cumulative lengths
+  int q_seq_start = cu_seqlens_q[batch_idx];
+  int q_seq_end = cu_seqlens_q[batch_idx + 1];
+  int k_seq_start = cu_seqlens_k[batch_idx];
+  int k_seq_end = cu_seqlens_k[batch_idx + 1];
+  
+  int q_seq_len = q_seq_end - q_seq_start;
+  int k_seq_len = k_seq_end - k_seq_start;
+  
+  // Check if this block is within the sequence
+  if (block_idx * BQ >= q_seq_len) {
+    return;
+  }
+  
+  // Calculate offsets in the packed tensor format
+  // Q/O shape: [total_tokens, num_heads, head_dim]
+  // K/V shape: [total_tokens, num_heads_kv, head_dim]
+  int q_offset = q_seq_start + block_idx * BQ;
+  int k_offset = k_seq_start;
+  
+  ulong kv_head_idx = head_idx / params->gqa_factor;
+  
+  // Move pointers to the correct position in packed format
+  Q += q_offset * params->H * params->D + head_idx * params->D;
+  K += k_offset * (params->H / params->gqa_factor) * params->D + kv_head_idx * params->D;
+  V += k_offset * (params->H / params->gqa_factor) * params->D + kv_head_idx * params->D;
+  O += q_offset * params->H * params->D + head_idx * params->D;
+  
+  if (has_mask) {
+    // Mask indexing would need to be updated based on the mask format
+    mask += batch_idx * mask_params->M_strides[0] + 
+            head_idx * mask_params->M_strides[1];
+  }
+
+  // Prepare threadgroup memory
+  constexpr short padQ = 16 / sizeof(T);
+  constexpr short padK = 16 / sizeof(T);
+  constexpr short padV = 16 / sizeof(T);
+
+  constexpr short LDQ_tgp = BD + padQ;
+  constexpr short LDK_tgp = BK + padK;
+  constexpr short LDV_tgp = BD + padV;
+
+  constexpr short tgp_mem_0 = (BK + padK) * (BD);
+  constexpr short tgp_mem_1 = BK * (BD + padV);
+  constexpr short tgp_mem_s = tgp_mem_0 > tgp_mem_1 ? tgp_mem_0 : tgp_mem_1;
+
+  threadgroup T Q_smem[BQ * (BD + padQ)];
+  threadgroup T KV_smem[tgp_mem_s];
+
+  threadgroup T* Qs = Q_smem;
+  threadgroup T* Ks = KV_smem;
+  threadgroup T* Vs = KV_smem;
+
+  // Prepare block loaders
+  using QBlockLoader = BlockLoaderT<
+      /* typename T = */ T,
+      /* short BROWS = */ BQ,
+      /* short BCOLS = */ BD,
+      /* short kDstStrRow = */ LDQ_tgp,
+      /* short kDstStrCol = */ 1,
+      /* short reduction_dim = */ 1,
+      /* short tgp_size = */ WM * WN * 32>;
+
+  // K is loaded in transposed
+  using KBlockLoader = BlockLoaderT<
+      /* typename T = */ T,
+      /* short BROWS = */ BK,
+      /* short BCOLS = */ BD,
+      /* short kDstStrRow = */ 1,
+      /* short kDstStrCol = */ LDK_tgp,
+      /* short reduction_dim = */ 0,
+      /* short tgp_size = */ WM * WN * 32>;
+
+  using VBlockLoader = BlockLoaderT<
+      /* typename T = */ T,
+      /* short BROWS = */ BK,
+      /* short BCOLS = */ BD,
+      /* short kDstStrRow = */ LDV_tgp,
+      /* short kDstStrCol = */ 1,
+      /* short reduction_dim = */ 0,
+      /* short tgp_size = */ WM * WN * 32>;
+
+  // For packed tensors, stride between tokens is H * D
+  int q_stride = params->H * params->D;
+  int kv_stride = (params->H / params->gqa_factor) * params->D;
+  
+  QBlockLoader loader_q(
+      Q, q_stride, Qs, simd_group_id, simd_lane_id);
+  KBlockLoader loader_k(
+      K, kv_stride, Ks, simd_group_id, simd_lane_id);
+  VBlockLoader loader_v(
+      V, kv_stride, Vs, simd_group_id, simd_lane_id);
+
+  // Apply softcapping adjustment to scale if needed
+  float adjusted_scale = params->scale;
+  if (params->softcapping != 1.0f) {
+    adjusted_scale = params->scale / params->softcapping;
+  }
+  TransformScale<T> ts(static_cast<T>(adjusted_scale * 1.44269504089));
+
+  // Prepare MMA tiles
+  constexpr short kFragSize = 8; // MMAFrag size
+  using MMAFrag_acc_t = BaseMMAFrag<AccumType, kFragSize, kFragSize>;
+
+  constexpr int kNWarps = WM * WN;
+  static_assert(
+      BQ >= (kNWarps * kFragSize) && BQ % (kNWarps * kFragSize) == 0,
+      "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence.");
+
+  // Q seq frags per warp
+  constexpr int TQ = BQ / (kNWarps * kFragSize);
+  // KV sequence frags (all warps load the same frags)
+  constexpr int TK = BK / kFragSize;
+  // HeadDim frags (all warps load the same frags)
+  constexpr int TD = BD / kFragSize;
+
+  static_assert(TQ == 1, "Check TQ");
+
+  MMATile<AccumType, TQ, 1, MMAFrag_acc_t> Qtile;
+  MMATile<AccumType, 1, TK, MMAFrag_acc_t> Ktile;
+  MMATile<AccumType, TQ, TK, MMAFrag_acc_t> Stile;
+  MMATile<AccumType, 1, 1, MMAFrag_acc_t> Vtile;
+  MMATile<AccumType, TQ, TD, MMAFrag_acc_t> Otile;
+
+  Otile.clear();
+
+  // Prepare mma tile offsets
+  const short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id);
+  const short sm = simd_coord.y;
+  const short sn = simd_coord.x;
+  const short tm = kFragSize * TQ * simd_group_id;
+
+  const short Qs_offset = (tm + sm) * LDQ_tgp + sn;
+  const short Ks_offset = sm * LDK_tgp + sn;
+  const short Vs_offset = sm * LDV_tgp + sn;
+
+  constexpr short Qs_tile_stride = kFragSize;
+  constexpr short Ks_tile_stride = kFragSize * LDK_tgp;
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Load Q blocks apply scale
+  int q_block_end = min(block_idx * BQ + BQ, q_seq_len);
+  int q_block_size = q_block_end - block_idx * BQ;
+  
+  if (q_block_size < BQ) {
+    loader_q.load_safe(short2(BD, q_block_size));
+  } else {
+    loader_q.load_unsafe();
+  }
+  loader_q.apply_inplace_op(ts);
+
+  // Init row reduction variables
+  constexpr short kRowsPT = decltype(Stile)::kRowsPerThread;
+
+  AccumType max_score[kRowsPT];
+  AccumType sum_score[kRowsPT] = {0};
+
+  // Init to -Inf
+  STEEL_PRAGMA_UNROLL
+  for (short i = 0; i < kRowsPT; ++i) {
+    max_score[i] = Limits<AccumType>::min;
+  }
+
+  // Calculate number of K blocks for this sequence.
+  // In general, we want to iterate over all key blocks.  However,
+  // when causal masking is enabled we only need to process up to the
+  // last key that influences this query block.  In decode mode
+  // (q_seq_len < k_seq_len), the single query token logically sits
+  // at the end of the key sequence.  Without adjusting for this the
+  // causal computation would incorrectly restrict processing to only
+  // the first key block, because the query position would appear to
+  // be at index 0.  To handle this we compute a causal_offset that
+  // shifts the query indices so they align with the end of the key
+  // sequence when q_seq_len < k_seq_len.
+  int kb_lim = (k_seq_len + BK - 1) / BK;
+
+  if (do_causal) {
+    // Offset the row indices for causal masking when the query length
+    // is smaller than the key length (decode mode).  This ensures
+    // that the computed row positions correspond to the correct
+    // positions within the key sequence.
+    int causal_offset = 0;
+    if (q_seq_len < k_seq_len) {
+      causal_offset = k_seq_len - q_seq_len;
+    }
+
+    // Determine the start/end of the current query block in the
+    // (possibly offset) sequence.  The block index operates on
+    // query positions but causal_offset places it relative to the
+    // key positions when in decode mode.
+    int q_block_start_in_seq = block_idx * BQ + causal_offset;
+    int q_block_end_in_seq = q_block_start_in_seq + q_block_size;
+
+    // Limit the number of key blocks so that blocks that are strictly
+    // beyond the last valid key (for this row) are not processed.
+    // When causal_offset > 0 this prevents prematurely exiting after
+    // the first block in decode mode.
+    kb_lim = min(kb_lim, (q_block_end_in_seq + BK - 1) / BK);
+  }
+
+  // Loop over KV seq length
+  for (int kb = 0; kb < kb_lim; kb++) {
+    // Load K block and apply scale
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    
+    int k_block_end = min(kb * BK + BK, k_seq_len);
+    int k_block_size = k_block_end - kb * BK;
+    
+    if (k_block_size < BK) {
+      loader_k.load_safe(short2(BD, k_block_size));
+    } else {
+      loader_k.load_unsafe();
+    }
+
+    // Do S = Q @ K.T
+    Stile.clear();
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    STEEL_PRAGMA_UNROLL
+    for (short dd = 0; dd < TD; dd++) {
+      simdgroup_barrier(mem_flags::mem_none);
+
+      Qtile.template load<T, 1, 1, LDQ_tgp, 1>(
+          &Qs[Qs_offset + dd * Qs_tile_stride]);
+      Ktile.template load<T, 1, 1, LDK_tgp, 1>(
+          &Ks[Ks_offset + dd * Ks_tile_stride]);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      tile_matmad(Stile, Qtile, Ktile, Stile);
+    }
+
+    // Mask out length sequence
+    if (k_block_size < BK) {
+      using stile_t = decltype(Stile);
+      using selem_t = typename stile_t::elem_type;
+      constexpr auto neg_inf = -metal::numeric_limits<selem_t>::infinity();
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < stile_t::kTileRows; i++) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < stile_t::kTileCols; j++) {
+          short col_pos = sn + (j * stile_t::kFragCols);
+          STEEL_PRAGMA_UNROLL
+          for (short jj = 0; jj < stile_t::MMAFrag_t::kElemCols; jj++) {
+            if ((col_pos + jj) >= k_block_size) {
+              Stile.frag_at(i, j)[jj] = neg_inf;
+            }
+          }
+        }
+      }
+    }
+
+    // Mask out if causal
+    if (do_causal) {
+      using stile_t = decltype(Stile);
+      using selem_t = typename stile_t::elem_type;
+      constexpr auto neg_inf = -metal::numeric_limits<selem_t>::infinity();
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < stile_t::kTileRows; i++) {
+        // Compute row position for causal mask.  In decode mode
+        // (q_seq_len < k_seq_len) the single query row should be
+        // aligned with the end of the key sequence.  Without this
+        // offset the row index would be zero and all but the first
+        // key block would be erroneously masked out.
+        int row_pos_causal = block_idx * BQ + tm + sm + (i * stile_t::kFragRows);
+        if (q_seq_len < k_seq_len) {
+          row_pos_causal += (k_seq_len - q_seq_len);
+        }
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < stile_t::kTileCols; j++) {
+          const int col_pos_in_seq = kb * BK + sn + (j * stile_t::kFragCols);
+          STEEL_PRAGMA_UNROLL
+          for (short jj = 0; jj < stile_t::MMAFrag_t::kElemCols; jj++) {
+            if (row_pos_causal < (col_pos_in_seq + jj)) {
+              Stile.frag_at(i, j)[jj] = neg_inf;
+            }
+          }
+        }
+      }
+    }
+
+    // Other masking as needed
+    if (has_mask) {
+      using stile_t = decltype(Stile);
+      using selem_t = typename stile_t::elem_type;
+      constexpr auto neg_inf = -metal::numeric_limits<selem_t>::infinity();
+
+      constexpr bool is_bool = is_same_v<MaskType, bool>;
+      using melem_t = typename metal::conditional_t<is_bool, bool, selem_t>;
+
+      using MMAFrag_mask_t = BaseMMAFrag<melem_t, kFragSize, kFragSize>;
+      using frag_t = typename MMAFrag_mask_t::frag_type;
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < stile_t::kTileRows; i++) {
+        // Use sequence-local positions
+        const int row_pos_in_seq = block_idx * BQ + tm + sm + (i * stile_t::kFragRows);
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < stile_t::kTileCols; j++) {
+          const int col_pos_in_seq = kb * BK + sn + (j * stile_t::kFragCols);
+
+          frag_t mfrag;
+
+          MMAFrag_mask_t::load_safe(
+              mfrag,
+              mask,
+              int(mask_params->M_strides[2]),
+              Int<1>{},
+              q_seq_len,
+              k_seq_len,
+              row_pos_in_seq,  // Already sequence-local
+              col_pos_in_seq); // Already sequence-local
+
+          STEEL_PRAGMA_UNROLL
+          for (short jj = 0; jj < stile_t::MMAFrag_t::kElemsPerFrag; jj++) {
+            if constexpr (is_bool) {
+              Stile.frag_at(i, j)[jj] =
+                  mfrag[jj] ? Stile.frag_at(i, j)[jj] : neg_inf;
+            } else {
+              Stile.frag_at(i, j)[jj] += selem_t(mfrag[jj]);
+            }
+          }
+        }
+      }
+    }
+
+    // Apply softcapping if needed (tanh(score) * softcapping)
+    if (params->softcapping != 1.0f) {
+      using stile_t = decltype(Stile);
+      using selem_t = typename stile_t::elem_type;
+      const selem_t softcapping_val = static_cast<selem_t>(params->softcapping);
+      
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < stile_t::kTileRows; i++) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < stile_t::kTileCols; j++) {
+          STEEL_PRAGMA_UNROLL
+          for (short jj = 0; jj < stile_t::MMAFrag_t::kElemsPerFrag; jj++) {
+            Stile.frag_at(i, j)[jj] = metal::tanh(Stile.frag_at(i, j)[jj]) * softcapping_val;
+          }
+        }
+      }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Load V blocks
+    if (k_block_size < BK) {
+      loader_v.load_safe(short2(BD, k_block_size));
+    } else {
+      loader_v.load_unsafe();
+    }
+
+    // Do softmax
+
+    // Temp variables
+    AccumType new_max[kRowsPT];
+    AccumType factor[kRowsPT];
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      new_max[i] = max_score[i];
+    }
+
+    // Row max
+    Stile.template row_reduce<MaxOp>(new_max);
+
+    // exp(Si - rowmax(Si))
+    Stile.template row_bin_op<ExpSubOp>(new_max);
+
+    // Factor exp(rowmax(Si) - rowmax(Si-1))
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      factor[i] = fast::exp2(max_score[i] - new_max[i]);
+    }
+
+    // Save max for next iteration
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      max_score[i] = new_max[i];
+    }
+
+    // Row Sum
+    AccumType sum_score_tmp[kRowsPT] = {0};
+    Stile.template row_reduce<SumOp>(sum_score_tmp);
+
+    // Update norm
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      sum_score[i] = sum_score[i] * factor[i] + sum_score_tmp[i];
+    }
+
+    // Update O
+    Otile.template row_bin_op<MulOp>(factor);
+
+    // Load V into registers
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    STEEL_PRAGMA_UNROLL
+    for (short iq = 0; iq < TQ; iq++) {
+      STEEL_PRAGMA_UNROLL
+      for (short id = 0; id < TD; id++) {
+        STEEL_PRAGMA_UNROLL
+        for (short ik = 0; ik < TK; ik++) {
+          if constexpr (BD == 128) {
+            simdgroup_barrier(mem_flags::mem_none);
+          }
+
+          const short kk = ik * kFragSize;
+          const short dd = id * kFragSize;
+
+          Vtile.template load<T, 1, 1, LDV_tgp, 1>(
+              &Vs[Vs_offset + kk * LDV_tgp + dd]);
+
+          if constexpr (BD == 128) {
+            simdgroup_barrier(mem_flags::mem_none);
+          }
+
+          MMAFrag_acc_t::mma(
+              Otile.frag_at(iq, id),
+              Stile.frag_at(iq, ik),
+              Vtile.frag_at(0, 0),
+              Otile.frag_at(iq, id));
+        }
+      }
+    }
+
+    // Prepare for next iteration
+    loader_k.next();
+    loader_v.next();
+  }
+
+  // Normalize output
+  Otile.template row_bin_op<DivOp>(sum_score);
+  threadgroup_barrier(mem_flags::mem_none);
+
+  // Store results
+  // O is already pointing to the correct block position from earlier adjustment
+  // Just need to offset within the block for this thread's tile
+  device T* O_tile = O + (tm + sm) * params->H * params->D + sn;
+
+  if (q_block_size < BQ) {
+    // Only store if this thread's tile is within the valid range
+    if ((tm + sm) < q_block_size && sn < BD) {
+      auto dst_tile_dims = short2(BD - sn, q_block_size - (tm + sm));
+      Otile.template store_safe<T, 1, 1>(O_tile, params->H * params->D, dst_tile_dims);
+    }
+  } else {
+    Otile.template store<T, 1, 1>(O_tile, params->H * params->D);
+  }
+}
+
+// clang-format off
+
+// SDPA full instantiations
+
+// Instantiate a templated kernel.
+// Extra args are used as template parameters:
+// e.g. instantiate_kernel(binary_int, binary, a, b) ->
+// [[host_name(binary_int)]] [kernel] binary<a, b>
+#define instantiate_kernel(name, func, ...) \
+  template [[host_name(                     \
+      name)]] [[kernel]] decltype(func<__VA_ARGS__>) func<__VA_ARGS__>;
+
+#define instantiate_attn(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \
+  instantiate_kernel(                                                    \
+      "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd            \
+      "_wm" #wm "_wn" #wn "_mask" #mname,                                \
+  attention, dtype, bq, bk, bd, wm, wn, mtype, float)
+
+#define instantiate_attn_shapes_helper(iname, itype, mname, mtype)  \
+    instantiate_attn(iname, itype, 16, 8, 256, 2, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 32, 32,  96, 4, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 32, 32,  80, 4, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 32, 32,  72, 4, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 32, 32,  64, 4, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 32, 32,  32, 4, 1, mname, mtype)
+
+#define instantiate_attn_mask_helper(iname, itype) \
+    instantiate_attn_shapes_helper(iname, itype, iname, itype) \
+    instantiate_attn_shapes_helper(iname, itype, bool_, bool)
+
+instantiate_attn_mask_helper(float16, half);
+instantiate_attn_mask_helper(bfloat16, bfloat16_t);
+instantiate_attn_mask_helper(float32, float);
+

sdpa-metal/scaled_dot_product_attention.mm

@@ -0,0 +1,316 @@
+#include <ATen/mps/MPSDevice.h>
+#include <ATen/mps/MPSStream.h>
+#include <torch/torch.h>
+
+// Include the auto-generated header with embedded metallib
+#ifdef EMBEDDED_METALLIB_HEADER
+#include EMBEDDED_METALLIB_HEADER
+#else
+#error "EMBEDDED_METALLIB_HEADER not defined"
+#endif
+
+static inline id<MTLBuffer> getMTLBufferStorage(const torch::Tensor &tensor) {
+  return __builtin_bit_cast(id<MTLBuffer>, tensor.storage().data());
+}
+
+// Helper function to get dtype string
+static std::string getDtypeString(torch::ScalarType dtype) {
+  switch (dtype) {
+  case torch::kFloat:
+    return "float32";
+  case torch::kHalf:
+    return "float16";
+  case torch::kBFloat16:
+    return "bfloat16";
+  default:
+    TORCH_CHECK(false, "Unsupported dtype for SDPA: ", dtype);
+  }
+}
+
+// Helper function to get dtype string for kernel names
+static std::string getKernelDtypeString(torch::ScalarType dtype) {
+  switch (dtype) {
+  case torch::kFloat:
+    return "float32";  // Match the instantiation names
+  case torch::kHalf:
+    return "float16";
+  case torch::kBFloat16:
+    return "bfloat16";
+  default:
+    TORCH_CHECK(false, "Unsupported dtype for SDPA: ", dtype);
+  }
+}
+
+
+// Parameters structure matching Flash Attention's AttnParams
+struct AttnParams {
+  int32_t B;              // batch size
+  int32_t H;              // number of heads
+  int32_t D;              // head dimension
+  int32_t qL;             // query sequence length (per sequence)
+  int32_t kL;             // key sequence length (per sequence)
+  int32_t gqa_factor;     // grouped query attention factor
+  float scale;            // attention scale
+  float softcapping;      // softcapping value (1.0 for no softcapping)
+  int32_t NQ;             // number of query blocks
+  int32_t NK;             // number of key blocks
+  int32_t NQ_aligned;     // aligned query blocks
+  int32_t NK_aligned;     // aligned key blocks
+  int32_t qL_rem;         // remainder query length
+  int32_t kL_rem;         // remainder key length
+  int32_t qL_off;         // query offset
+  int64_t Q_strides[3];   // query tensor strides
+  int64_t K_strides[3];   // key tensor strides
+  int64_t V_strides[3];   // value tensor strides
+  int64_t O_strides[3];   // output tensor strides
+  
+  // Flash Attention variable-length support
+  int32_t total_q_tokens; // Total number of query tokens
+  int32_t total_k_tokens; // Total number of key/value tokens
+  int32_t max_seqlen_q;   // Maximum query sequence length
+  int32_t max_seqlen_k;   // Maximum key/value sequence length
+};
+
+// Forward declarations for kernel implementations
+void call_flash_attention_varlen(
+    id<MTLDevice> device,
+    id<MTLCommandBuffer> cmdBuf,
+    id<MTLLibrary> lib,
+    torch::Tensor &out,
+    torch::Tensor &query,
+    torch::Tensor &key,
+    torch::Tensor &value,
+    torch::Tensor &cu_seqlens_q,
+    torch::Tensor &cu_seqlens_k,
+    int64_t max_seqlen_q,
+    int64_t max_seqlen_k,
+    bool do_causal,
+    double scale,
+    double softcapping);
+
+
+void flash_attention_varlen(
+    torch::Tensor &out,           // [total_q_tokens, num_heads, head_size]
+    torch::Tensor &query,         // [total_q_tokens, num_heads, head_size]
+    torch::Tensor &key,           // [total_k_tokens, num_heads_kv, head_size]
+    torch::Tensor &value,         // [total_k_tokens, num_heads_kv, head_size]
+    torch::Tensor &cu_seqlens_q,  // [batch_size + 1]
+    torch::Tensor &cu_seqlens_k,  // [batch_size + 1]
+    int64_t max_seqlen_q,         // Maximum query sequence length
+    int64_t max_seqlen_k,         // Maximum key sequence length
+    bool do_causal,               // Whether to use causal mask
+    double scale,                 // Attention scale
+    double softcapping) {         // Softcapping value
+
+  try {
+    // Get device and stream
+    id<MTLDevice> device = at::mps::MPSDevice::getInstance()->device();
+    at::mps::MPSStream *stream = at::mps::getCurrentMPSStream();
+    TORCH_CHECK(stream, "Failed to get current MPS stream");
+
+  // Get dimensions from Flash Attention format
+  int64_t total_q_tokens = query.size(0);
+  int64_t num_heads = query.size(1);
+  int64_t head_dim = query.size(2);
+  int64_t num_heads_kv = key.size(1);
+  int64_t batch_size = cu_seqlens_q.size(0) - 1;  // cu_seqlens has batch_size + 1 elements
+  
+  // Check if we support this head dimension
+  std::vector<int> supported_head_dims = {32, 64, 72, 80, 96, 128, 256};
+  bool supported_head_dim = std::find(supported_head_dims.begin(), 
+                                      supported_head_dims.end(), 
+                                      head_dim) != supported_head_dims.end();
+  
+  TORCH_CHECK(supported_head_dim, "Head dimension ", head_dim, " is not supported");
+  TORCH_CHECK(cu_seqlens_q.size(0) == cu_seqlens_k.size(0), 
+              "cu_seqlens_q and cu_seqlens_k must have the same size");
+
+  // Load Metal library
+  static id<MTLLibrary> lib = nil;
+  if (!lib) {
+    NSError *error = nil;
+    lib = EMBEDDED_METALLIB_NAMESPACE::createLibrary(device, &error);
+    TORCH_CHECK(lib,
+                "Failed to create Metal library from embedded data: ",
+                error.localizedDescription.UTF8String);
+  }
+
+  // Get command buffer
+  id<MTLCommandBuffer> cmdBuf = stream->commandBuffer();
+  TORCH_CHECK(cmdBuf, "Failed to get MPS command buffer");
+
+  // For variable-length Flash Attention, always use the full attention kernel
+  
+  // Call the Flash Attention kernel
+  call_flash_attention_varlen(device, cmdBuf, lib, out, query, key, value, 
+                              cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
+                              do_causal, scale, softcapping);
+  } catch (const std::exception& e) {
+    throw;
+  } catch (...) {
+    throw;
+  }
+}
+
+// Implementation of Flash Attention variable-length kernel
+void call_flash_attention_varlen(
+    id<MTLDevice> device,
+    id<MTLCommandBuffer> cmdBuf,
+    id<MTLLibrary> lib,
+    torch::Tensor &out,
+    torch::Tensor &query,
+    torch::Tensor &key,
+    torch::Tensor &value,
+    torch::Tensor &cu_seqlens_q,
+    torch::Tensor &cu_seqlens_k,
+    int64_t max_seqlen_q,
+    int64_t max_seqlen_k,
+    bool do_causal,
+    double scale,
+    double softcapping) {
+
+  // Get dimensions
+  int64_t total_q_tokens = query.size(0);
+  int64_t num_heads = query.size(1);
+  int64_t head_dim = query.size(2);
+  int64_t num_heads_kv = key.size(1);
+  int64_t batch_size = cu_seqlens_q.size(0) - 1;
+
+  // Grouped Query Attention factor
+  int32_t gqa_factor = num_heads / num_heads_kv;
+
+  // Block sizes based on head dimension
+  const int BQ = (head_dim == 256) ? 16 : 32;  // Use BQ=16 for head_dim=256
+  const int bk = (head_dim == 256) ? 8 : ((head_dim >= 128) ? 16 : 32);  // Use bk=8 for head_dim=256
+  const int WM = (head_dim == 256) ? 2 : 4;  // Use WM=2 for head_dim=256
+  const int WN = 1;
+
+  // Setup parameters
+  AttnParams params = {}; // Zero-initialize all fields
+  params.B = batch_size;
+  params.H = num_heads;
+  params.D = head_dim;
+  params.gqa_factor = gqa_factor;
+  params.scale = static_cast<float>(scale);
+  params.softcapping = static_cast<float>(softcapping);
+  params.total_q_tokens = total_q_tokens;
+  params.total_k_tokens = key.size(0);
+  params.max_seqlen_q = max_seqlen_q;
+  params.max_seqlen_k = max_seqlen_k;
+  
+  // Initialize fields that might be checked but aren't used in Flash Attention
+  params.qL = 0;  // Not used in variable-length attention
+  params.kL = 0;  // Not used in variable-length attention
+  params.NQ = 0;  // Not used
+  params.NK = 0;  // Not used
+  params.NQ_aligned = 0;
+  params.NK_aligned = 0;
+  params.qL_rem = 0;
+  params.kL_rem = 0;
+  params.qL_off = 0;
+  
+  // Strides are not used for packed tensors (contiguous)
+  params.Q_strides[0] = 0;
+  params.Q_strides[1] = 0;
+  params.Q_strides[2] = 0;
+  params.K_strides[0] = 0;
+  params.K_strides[1] = 0;
+  params.K_strides[2] = 0;
+  params.V_strides[0] = 0;
+  params.V_strides[1] = 0;
+  params.V_strides[2] = 0;
+  params.O_strides[0] = 0;
+  params.O_strides[1] = 0;
+  params.O_strides[2] = 0;
+  
+  // For variable-length attention, we'll process each sequence separately
+  // The kernel will handle the cu_seqlens internally
+
+  bool has_mask = false;  // Masks are not supported in Flash Attention
+
+  // Setup function constants
+  MTLFunctionConstantValues *constants = [MTLFunctionConstantValues new];
+  [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:300];
+  [constants setConstantValue:&do_causal type:MTLDataTypeBool atIndex:301];
+
+  // Construct kernel name based on data type and head dimension
+  std::string kernel_name = "steel_attention_";
+  kernel_name += getKernelDtypeString(query.scalar_type());
+  kernel_name += "_bq" + std::to_string(BQ);
+  kernel_name += "_bk" + std::to_string(bk);
+  kernel_name += "_bd" + std::to_string(head_dim);
+  kernel_name += "_wm" + std::to_string(WM) + "_wn" + std::to_string(WN);
+  kernel_name += "_maskbool_";  // Always use bool for mask type (no masks supported)
+
+  // Get kernel function
+  NSError *error = nil;
+  id<MTLFunction> function = [lib newFunctionWithName:[NSString stringWithUTF8String:kernel_name.c_str()]
+                                      constantValues:constants
+                                              error:&error];
+  TORCH_CHECK(function, "Failed to get Metal function: ", kernel_name,
+              " Error: ", error ? error.localizedDescription.UTF8String : "unknown");
+
+  // Create compute pipeline
+  id<MTLComputePipelineState> pipeline = [device newComputePipelineStateWithFunction:function error:&error];
+  TORCH_CHECK(pipeline, "Failed to create compute pipeline: ",
+              error ? error.localizedDescription.UTF8String : "unknown");
+
+  // Setup command encoder with dispatch sync
+  at::mps::MPSStream *stream = at::mps::getCurrentMPSStream();
+  dispatch_queue_t q = stream->queue();
+  dispatch_sync(q, ^{
+    id<MTLComputeCommandEncoder> encoder = [cmdBuf computeCommandEncoder];
+    TORCH_CHECK(encoder, "Failed to create compute encoder");
+
+    [encoder setComputePipelineState:pipeline];
+    
+    // Set buffers
+    int buffer_idx = 0;
+    
+    // Query buffer - index 0
+    [encoder setBuffer:getMTLBufferStorage(query) 
+              offset:query.storage_offset() * query.element_size() 
+              atIndex:buffer_idx++];
+    
+    // Key buffer - index 1
+    [encoder setBuffer:getMTLBufferStorage(key) 
+              offset:key.storage_offset() * key.element_size() 
+              atIndex:buffer_idx++];
+    
+    // Value buffer - index 2
+    [encoder setBuffer:getMTLBufferStorage(value) 
+              offset:value.storage_offset() * value.element_size() 
+              atIndex:buffer_idx++];
+    
+    // Output buffer - index 3
+    [encoder setBuffer:getMTLBufferStorage(out) 
+              offset:out.storage_offset() * out.element_size() 
+              atIndex:buffer_idx++];
+    
+    // Parameters - index 4
+    [encoder setBytes:&params length:sizeof(AttnParams) atIndex:buffer_idx++];
+    
+    // Skip mask parameters - indices 5 and 6 (masks not supported)
+    buffer_idx += 2;
+    
+    // Set cu_seqlens buffers - indices 7 and 8
+    [encoder setBuffer:getMTLBufferStorage(cu_seqlens_q) 
+              offset:cu_seqlens_q.storage_offset() * cu_seqlens_q.element_size() 
+              atIndex:7];
+    [encoder setBuffer:getMTLBufferStorage(cu_seqlens_k) 
+              offset:cu_seqlens_k.storage_offset() * cu_seqlens_k.element_size() 
+              atIndex:8];
+
+    // Calculate grid dimensions
+    // We need to process each sequence independently
+    int64_t max_blocks_q = (max_seqlen_q + BQ - 1) / BQ;
+    
+    MTLSize gridSize = MTLSizeMake(max_blocks_q, num_heads, batch_size);
+    MTLSize threadgroupSize = MTLSizeMake(32, WM, WN);
+    
+    [encoder dispatchThreadgroups:gridSize threadsPerThreadgroup:threadgroupSize];
+    [encoder endEncoding];
+    
+    stream->synchronize(at::mps::SyncType::COMMIT);
+  });
+}

tests/test_flash_attention.py

@@ -0,0 +1,480 @@
+import torch
+import pytest
+import metal_flash_sdpa
+
+
+def create_cu_seqlens(seq_lengths):
+    """Create cumulative sequence lengths tensor."""
+    cu_seqlens = [0]
+    for length in seq_lengths:
+        cu_seqlens.append(cu_seqlens[-1] + length)
+    return torch.tensor(cu_seqlens, dtype=torch.int32, device="mps")
+
+
+def compute_attention_reference(query, key, value, scale, causal=False, softcapping=1.0, gqa_ratio=1):
+    """Compute reference attention output for validation."""
+    num_heads = query.shape[1]
+    expected = torch.zeros_like(query)
+    
+    for h in range(num_heads):
+        kv_h = h // gqa_ratio if gqa_ratio > 1 else h
+        q_h = query[:, h, :]
+        k_h = key[:, kv_h, :]
+        v_h = value[:, kv_h, :]
+        
+        scores = torch.matmul(q_h, k_h.transpose(-1, -2)) * scale
+        
+        # Apply softcapping if not 1.0
+        if softcapping != 1.0:
+            scores = scores / softcapping
+            scores = torch.tanh(scores) * softcapping
+        
+        # Apply causal mask if needed
+        if causal:
+            seq_len = query.shape[0]
+            causal_mask = torch.triu(torch.ones(seq_len, seq_len, device="mps"), diagonal=1).bool()
+            scores.masked_fill_(causal_mask, float("-inf"))
+        
+        attn_weights = torch.softmax(scores, dim=-1)
+        expected[:, h, :] = torch.matmul(attn_weights, v_h)
+    
+    return expected
+
+
+def get_tolerance(dtype, head_dim):
+    """Get appropriate tolerance based on dtype and head dimension."""
+    if dtype == torch.bfloat16:
+        return (2e-2, 2e-2) if head_dim >= 96 else (1.6e-2, 1.6e-2)
+    elif dtype == torch.float16:
+        return (2e-3, 2e-3)
+    else:
+        return (1e-3, 1e-3)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("head_dim", [32, 64, 72, 80, 96, 128, 256])
+@pytest.mark.parametrize("seq_config", [
+    # (seq_lengths_q, seq_lengths_k, description)
+    ([32], [32], "single_sequence"),
+    ([8, 16, 12], [10, 20, 15], "variable_lengths"),
+    ([16, 24], [16, 24], "multiple_sequences"),
+    ([2], [2], "small_sequence_2"),
+    ([4], [4], "small_sequence_4"),
+    ([8], [8], "small_sequence_8"),
+    ([16], [32], "cross_attention_q_lt_k"),
+    ([32], [16], "cross_attention_q_gt_k"),
+    ([8], [128], "cross_attention_large_diff"),
+    ([1], [64], "single_query_token"),
+])
+@pytest.mark.parametrize("causal", [False, True])
+def test_flash_attention_comprehensive(dtype, head_dim, seq_config, causal):
+    """Comprehensive test for Flash Attention with various configurations."""
+    torch.manual_seed(42)
+    
+    seq_lengths_q, seq_lengths_k, _ = seq_config
+    
+    # Skip causal tests for cross-attention cases
+    if causal and seq_lengths_q != seq_lengths_k:
+        pytest.skip("Causal attention only valid when q_seq == k_seq")
+    
+    # Test parameters
+    num_heads = 4
+    
+    # Create cumulative sequence lengths
+    cu_seqlens_q = create_cu_seqlens(seq_lengths_q)
+    cu_seqlens_k = create_cu_seqlens(seq_lengths_k)
+    
+    total_q = sum(seq_lengths_q)
+    total_k = sum(seq_lengths_k)
+    max_seqlen_q = max(seq_lengths_q)
+    max_seqlen_k = max(seq_lengths_k)
+    
+    # Create input tensors
+    query = torch.randn(total_q, num_heads, head_dim, dtype=dtype, device="mps")
+    key = torch.randn(total_k, num_heads, head_dim, dtype=dtype, device="mps")
+    value = torch.randn(total_k, num_heads, head_dim, dtype=dtype, device="mps")
+    
+    # Scale factor
+    scale = 1.0 / (head_dim ** 0.5)
+    
+    # Call Flash Attention
+    out = torch.empty_like(query)
+    metal_flash_sdpa.flash_attention_varlen(
+        out=out,
+        query=query,
+        key=key,
+        value=value,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        do_causal=causal,
+        scale=scale,
+        softcapping=1.0,
+    )
+    
+    # Compute ground truth for each sequence
+    expected = torch.zeros_like(out)
+    batch_size = len(seq_lengths_q)
+    
+    for i in range(batch_size):
+        q_start, q_end = cu_seqlens_q[i].item(), cu_seqlens_q[i+1].item()
+        k_start, k_end = cu_seqlens_k[i].item(), cu_seqlens_k[i+1].item()
+        
+        if q_end > q_start and k_end > k_start:  # Skip empty sequences
+            q_i = query[q_start:q_end]
+            k_i = key[k_start:k_end]
+            v_i = value[k_start:k_end]
+            
+            expected_i = compute_attention_reference(q_i, k_i, v_i, scale, causal=causal)
+            expected[q_start:q_end] = expected_i
+    
+    # Check results
+    rtol, atol = get_tolerance(dtype, head_dim)
+    torch.testing.assert_close(out, expected, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("head_dim", [32, 64, 72, 80, 96, 128, 256])
+@pytest.mark.parametrize("gqa_config", [
+    # (num_heads, num_kv_heads, seq_len)
+    (8, 2, 32),    # 4:1 ratio
+    (16, 4, 32),   # 4:1 ratio
+    (16, 8, 32),   # 2:1 ratio
+    (16, 2, 32),   # 8:1 ratio
+    (16, 4, 128),  # 4:1 ratio with larger sequence
+])
+def test_flash_attention_gqa(dtype, head_dim, gqa_config):
+    """Test Flash Attention with Grouped Query Attention configurations."""
+    torch.manual_seed(42)
+    
+    num_heads, num_kv_heads, seq_len = gqa_config
+    gqa_ratio = num_heads // num_kv_heads
+    
+    # Create cumulative sequence lengths
+    cu_seqlens = create_cu_seqlens([seq_len])
+    
+    # Create input tensors
+    query = torch.randn(seq_len, num_heads, head_dim, dtype=dtype, device="mps")
+    key = torch.randn(seq_len, num_kv_heads, head_dim, dtype=dtype, device="mps")
+    value = torch.randn(seq_len, num_kv_heads, head_dim, dtype=dtype, device="mps")
+    
+    # Scale factor
+    scale = 1.0 / (head_dim ** 0.5)
+    
+    # Call Flash Attention
+    out = torch.empty_like(query)
+    metal_flash_sdpa.flash_attention_varlen(
+        out=out,
+        query=query,
+        key=key,
+        value=value,
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_k=cu_seqlens,
+        max_seqlen_q=seq_len,
+        max_seqlen_k=seq_len,
+        do_causal=False,
+        scale=scale,
+        softcapping=1.0,
+    )
+    
+    # Compute ground truth with GQA
+    expected = compute_attention_reference(query, key, value, scale, gqa_ratio=gqa_ratio)
+    
+    # Check results
+    rtol, atol = get_tolerance(dtype, head_dim)
+    torch.testing.assert_close(out, expected, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("softcapping_config", [
+    # (softcapping_value, seq_lengths, head_dim)
+    (1.0, [32], 64),        # No softcapping
+    (50.0, [32, 24], 64),   # Regular softcapping
+    (10.0, [16], 128),      # Strong softcapping
+    (1000.0, [16], 64),     # Very weak softcapping
+    (30.0, [48], 96),       # Medium softcapping
+])
+def test_flash_attention_softcapping(dtype, softcapping_config):
+    """Test Flash Attention with various softcapping values."""
+    torch.manual_seed(42)
+    
+    softcapping, seq_lengths, head_dim = softcapping_config
+    num_heads = 4
+    
+    # Create cumulative sequence lengths
+    cu_seqlens = create_cu_seqlens(seq_lengths)
+    total_tokens = sum(seq_lengths)
+    max_seqlen = max(seq_lengths)
+    
+    # Create input tensors
+    query = torch.randn(total_tokens, num_heads, head_dim, dtype=dtype, device="mps")
+    key = torch.randn(total_tokens, num_heads, head_dim, dtype=dtype, device="mps")
+    value = torch.randn(total_tokens, num_heads, head_dim, dtype=dtype, device="mps")
+    
+    # Scale factor
+    scale = 1.0 / (head_dim ** 0.5)
+    
+    # Call Flash Attention with softcapping
+    out = torch.empty_like(query)
+    metal_flash_sdpa.flash_attention_varlen(
+        out=out,
+        query=query,
+        key=key,
+        value=value,
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_k=cu_seqlens,
+        max_seqlen_q=max_seqlen,
+        max_seqlen_k=max_seqlen,
+        do_causal=False,
+        scale=scale,
+        softcapping=softcapping,
+    )
+    
+    # Compute ground truth with softcapping
+    expected = torch.zeros_like(query)
+    
+    for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:]):
+        if end > start:
+            q_seq = query[start:end]
+            k_seq = key[start:end]
+            v_seq = value[start:end]
+            
+            expected_seq = compute_attention_reference(
+                q_seq, k_seq, v_seq, scale, softcapping=softcapping
+            )
+            expected[start:end] = expected_seq
+    
+    # Check results (higher tolerance for softcapping)
+    # Note: Softcapping with strong values (< 50) has higher error due to 
+    # the interaction between tanh transformation and exp2-based softmax
+    if dtype == torch.bfloat16:
+        if softcapping < 50:
+            rtol, atol = 1.5e-1, 1.5e-1  # Higher tolerance for strong softcapping
+        else:
+            rtol, atol = 3e-2, 3e-2
+    elif dtype == torch.float16:
+        if softcapping < 50:
+            rtol, atol = 1e-1, 1e-1
+        else:
+            rtol, atol = 2e-2, 2e-2
+    else:
+        if softcapping < 50:
+            rtol, atol = 1.5e-1, 1.5e-1  # Higher tolerance for strong softcapping with float32
+        else:
+            rtol, atol = 1e-2, 1e-2
+    torch.testing.assert_close(out, expected, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("large_seq_config", [
+    # (q_seq, k_seq, head_dim, dtype)
+    (32, 2048, 64, torch.float32),
+    (16, 1024, 96, torch.float16),
+    (64, 1536, 64, torch.bfloat16),
+])
+def test_flash_attention_large_sequences(large_seq_config):
+    """Test Flash Attention with large k sequences (>= 1024)."""
+    torch.manual_seed(42)
+    
+    q_seq, k_seq, head_dim, dtype = large_seq_config
+    num_heads = 4
+    
+    # Create cumulative sequence lengths
+    cu_seqlens_q = create_cu_seqlens([q_seq])
+    cu_seqlens_k = create_cu_seqlens([k_seq])
+    
+    # Create input tensors
+    query = torch.randn(q_seq, num_heads, head_dim, dtype=dtype, device="mps")
+    key = torch.randn(k_seq, num_heads, head_dim, dtype=dtype, device="mps")
+    value = torch.randn(k_seq, num_heads, head_dim, dtype=dtype, device="mps")
+    
+    # Scale factor
+    scale = 1.0 / (head_dim ** 0.5)
+    
+    # Call Flash Attention
+    out = torch.empty_like(query)
+    metal_flash_sdpa.flash_attention_varlen(
+        out=out,
+        query=query,
+        key=key,
+        value=value,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=q_seq,
+        max_seqlen_k=k_seq,
+        do_causal=False,
+        scale=scale,
+        softcapping=1.0,
+    )
+    
+    # Compute ground truth
+    expected = compute_attention_reference(query, key, value, scale)
+    
+    # Check results (higher tolerance for large sequences)
+    if dtype == torch.bfloat16:
+        rtol, atol = 3e-2, 3e-2
+    elif dtype == torch.float16:
+        rtol, atol = 5e-3, 5e-3
+    else:
+        rtol, atol = 2e-3, 2e-3
+    torch.testing.assert_close(out, expected, rtol=rtol, atol=atol)
+
+
+def test_flash_attention_edge_cases():
+    """Test Flash Attention edge cases."""
+    torch.manual_seed(42)
+    
+    # Test 1: Single token sequence
+    query = torch.randn(1, 1, 64, device="mps")
+    key = torch.randn(1, 1, 64, device="mps")
+    value = torch.randn(1, 1, 64, device="mps")
+    cu_seqlens = create_cu_seqlens([1])
+    out = torch.empty_like(query)
+    
+    metal_flash_sdpa.flash_attention_varlen(
+        out=out,
+        query=query,
+        key=key,
+        value=value,
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_k=cu_seqlens,
+        max_seqlen_q=1,
+        max_seqlen_k=1,
+        do_causal=False,
+        scale=0.125,
+        softcapping=1.0,
+    )
+    
+    # With single token, output should equal value
+    torch.testing.assert_close(out, value, rtol=1e-5, atol=1e-5)
+    
+    # Test 2: Empty sequence in batch
+    seq_lengths = [8, 0, 12]  # Middle sequence is empty
+    cu_seqlens = create_cu_seqlens(seq_lengths)
+    total_tokens = sum(seq_lengths)
+    
+    query = torch.randn(total_tokens, 4, 64, device="mps")
+    key = torch.randn(total_tokens, 4, 64, device="mps")
+    value = torch.randn(total_tokens, 4, 64, device="mps")
+    out = torch.empty_like(query)
+    
+    # This should handle empty sequences gracefully
+    metal_flash_sdpa.flash_attention_varlen(
+        out=out,
+        query=query,
+        key=key,
+        value=value,
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_k=cu_seqlens,
+        max_seqlen_q=max(seq_lengths) if seq_lengths else 0,
+        max_seqlen_k=max(seq_lengths) if seq_lengths else 0,
+        do_causal=False,
+        scale=0.125,
+        softcapping=1.0,
+    )
+
+
+def test_flash_attention_unsupported_cases():
+    """Test that unsupported cases raise appropriate errors."""
+    
+    # Test 1: Unsupported head dimension
+    query = torch.randn(16, 4, 48, device="mps")  # head_dim = 48 (not supported)
+    key = torch.randn(16, 4, 48, device="mps")
+    value = torch.randn(16, 4, 48, device="mps")
+    cu_seqlens = create_cu_seqlens([16])
+    out = torch.empty_like(query)
+    
+    with pytest.raises(RuntimeError, match="Head dimension .* is not supported"):
+        metal_flash_sdpa.flash_attention_varlen(
+            out=out,
+            query=query,
+            key=key,
+            value=value,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=16,
+            max_seqlen_k=16,
+            do_causal=False,
+            scale=0.144,
+            softcapping=1.0,
+        )
+    
+    # Test 2: Wrong dtype for cu_seqlens (should be int32)
+    cu_seqlens_wrong = torch.tensor([0, 16], dtype=torch.int64, device="mps")
+    query = torch.randn(16, 4, 64, device="mps")
+    key = torch.randn(16, 4, 64, device="mps")
+    value = torch.randn(16, 4, 64, device="mps")
+    
+    # This will silently fail (output will be unchanged)
+    out = torch.full_like(query, -999.0)
+    metal_flash_sdpa.flash_attention_varlen(
+        out=out,
+        query=query,
+        key=key,
+        value=value,
+        cu_seqlens_q=cu_seqlens_wrong,
+        cu_seqlens_k=cu_seqlens_wrong,
+        max_seqlen_q=16,
+        max_seqlen_k=16,
+        do_causal=False,
+        scale=0.125,
+        softcapping=1.0,
+    )
+    
+    # Check that output wasn't modified (kernel didn't run)
+    assert (out == -999.0).all(), "cu_seqlens with wrong dtype should cause kernel to not run"
+
+
+def test_flash_attn_varlen_func():
+    """Test the flash_attn_varlen_func compatibility function."""
+    torch.manual_seed(42)
+    
+    # Test dimensions
+    seq_lengths = [8, 12]
+    num_heads = 4
+    head_dim = 64
+    
+    # Create cumulative sequence lengths
+    cu_seqlens = create_cu_seqlens(seq_lengths)
+    total_tokens = sum(seq_lengths)
+    max_seqlen = max(seq_lengths)
+    
+    # Create input tensors
+    q = torch.randn(total_tokens, num_heads, head_dim, device="mps")
+    k = torch.randn(total_tokens, num_heads, head_dim, device="mps")
+    v = torch.randn(total_tokens, num_heads, head_dim, device="mps")
+    
+    # Call the compatibility function
+    out = metal_flash_sdpa.flash_attn_varlen_func(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_k=cu_seqlens,
+        max_seqlen_q=max_seqlen,
+        max_seqlen_k=max_seqlen,
+        dropout_p=0.0,
+        softmax_scale=None,  # Will use 1/sqrt(head_dim)
+        causal=False,
+    )
+    
+    # Check that output has correct shape and is not zeros
+    assert out.shape == q.shape
+    assert out.abs().max().item() > 0
+    
+    # Test with causal
+    out_causal = metal_flash_sdpa.flash_attn_varlen_func(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_k=cu_seqlens,
+        max_seqlen_q=max_seqlen,
+        max_seqlen_k=max_seqlen,
+        dropout_p=0.0,
+        softmax_scale=0.125,
+        causal=True,
+    )
+    
+    assert out_causal.shape == q.shape
+    assert out_causal.abs().max().item() > 0

torch-ext/metal_flash_sdpa/__init__.py

@@ -0,0 +1,11 @@
+from ._custom_ops import (
+    flash_attention_varlen,
+    flash_attn_varlen_func,
+)
+from ._ops import ops
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+    "ops",
+]

torch-ext/metal_flash_sdpa/_custom_ops.py

@@ -0,0 +1,117 @@
+from typing import List, Optional
+
+import torch
+
+from ._ops import ops
+
+
+def flash_attention_varlen(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    do_causal: bool = False,
+    scale: Optional[float] = None,
+    softcapping: float = 1.0,
+) -> None:
+    """
+    Flash Attention with variable-length sequences.
+    
+    Args:
+        out: Output tensor of shape [total_q_tokens, num_heads, head_dim]
+        query: Query tensor of shape [total_q_tokens, num_heads, head_dim]
+        key: Key tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        value: Value tensor of shape [total_k_tokens, num_heads_kv, head_dim]
+        cu_seqlens_q: Cumulative sequence lengths for queries, shape [batch_size + 1], dtype must be torch.int32
+        cu_seqlens_k: Cumulative sequence lengths for keys, shape [batch_size + 1], dtype must be torch.int32
+        max_seqlen_q: Maximum sequence length in the query batch
+        max_seqlen_k: Maximum sequence length in the key batch
+        do_causal: Whether to apply causal masking
+        scale: Attention scale factor (default: 1/sqrt(head_dim))
+        softcapping: Softcapping value (default: 1.0, must be 1.0 for this implementation)
+    
+    Note:
+        - cu_seqlens_q and cu_seqlens_k must have dtype torch.int32 for Metal compatibility
+        - Supported head dimensions: 32, 64, 72, 80, 96, 128
+        - Masks are not supported
+    """
+    if scale is None:
+        scale = query.shape[-1] ** -0.5
+    
+    ops.flash_attention_varlen(
+        out,
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        do_causal,
+        scale,
+        softcapping,
+    )
+
+def flash_attn_varlen_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: tuple = (-1, -1),
+    alibi_slopes: Optional[torch.Tensor] = None,
+    deterministic: bool = False,
+    return_attn_probs: bool = False,
+) -> torch.Tensor:
+    """
+    Flash Attention function with API compatible with the original Flash Attention.
+    
+    Note: This implementation does not support:
+    - dropout
+    - window attention
+    - alibi slopes
+    - returning attention probabilities
+    """
+    if dropout_p > 0:
+        raise NotImplementedError("Dropout is not supported in this implementation")
+    if window_size != (-1, -1):
+        raise NotImplementedError("Window attention is not supported")
+    if alibi_slopes is not None:
+        raise NotImplementedError("ALiBi is not supported")
+    if return_attn_probs:
+        raise NotImplementedError("Returning attention probabilities is not supported")
+    
+    # Create output tensor
+    out = torch.empty_like(q)
+    
+    # Call the kernel
+    flash_attention_varlen(
+        out=out,
+        query=q,
+        key=k,
+        value=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        do_causal=causal,
+        scale=softmax_scale,
+        softcapping=1.0,
+    )
+    
+    return out
+
+
+__all__ = [
+    "flash_attention_varlen",
+    "flash_attn_varlen_func",
+]

torch-ext/torch_binding.cpp

@@ -0,0 +1,11 @@
+#include <torch/library.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("flash_attention_varlen(Tensor! out, Tensor query, Tensor key, Tensor value, Tensor cu_seqlens_q, Tensor cu_seqlens_k, int max_seqlen_q, int max_seqlen_k, bool do_causal, float scale, float softcapping) -> ()");
+  ops.impl("flash_attention_varlen", torch::kMPS, flash_attention_varlen);
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -0,0 +1,16 @@
+#pragma once
+
+#include <torch/torch.h>
+
+void flash_attention_varlen(
+    torch::Tensor &out,
+    torch::Tensor &query,
+    torch::Tensor &key,
+    torch::Tensor &value,
+    torch::Tensor &cu_seqlens_q,
+    torch::Tensor &cu_seqlens_k,
+    int64_t max_seqlen_q,
+    int64_t max_seqlen_k,
+    bool do_causal,
+    double scale,
+    double softcapping);