mps-flash-attn 0.1.7__cp314-cp314-macosx_15_0_arm64.whl

mps_flash_attn/__init__.py

@@ -0,0 +1,264 @@
+"""
+MPS Flash Attention - Flash Attention for PyTorch on Apple Silicon
+
+This package provides memory-efficient attention using Metal Flash Attention kernels.
+"""
+
+__version__ = "0.1.7"
+
+import torch
+from typing import Optional
+import math
+import threading
+import os
+
+# Try to import the C++ extension
+try:
+    from . import _C
+    _HAS_MFA = True
+except ImportError as e:
+    _HAS_MFA = False
+    _IMPORT_ERROR = str(e)
+
+# Note: The C++ extension handles loading libMFABridge.dylib via dlopen.
+# Set MFA_BRIDGE_PATH environment variable to specify the library location.
+# Do NOT load the library here via ctypes - that causes duplicate class warnings.
+
+
+def is_available() -> bool:
+    """Check if MPS Flash Attention is available."""
+    return _HAS_MFA and torch.backends.mps.is_available()
+
+
+class FlashAttentionFunction(torch.autograd.Function):
+    """Autograd function for Flash Attention with backward pass support."""
+
+    @staticmethod
+    def forward(ctx, query, key, value, is_causal, scale, attn_mask):
+        # Apply scale if provided (MFA uses 1/sqrt(D) internally)
+        scale_factor = 1.0
+        if scale is not None:
+            default_scale = 1.0 / math.sqrt(query.shape[-1])
+            if abs(scale - default_scale) > 1e-6:
+                scale_factor = scale / default_scale
+                query = query * scale_factor
+
+        # Forward with logsumexp for backward
+        output, logsumexp = _C.forward_with_lse(query, key, value, is_causal, attn_mask)
+
+        # Save for backward
+        if attn_mask is not None:
+            ctx.save_for_backward(query, key, value, output, logsumexp, attn_mask)
+            ctx.has_mask = True
+        else:
+            ctx.save_for_backward(query, key, value, output, logsumexp)
+            ctx.has_mask = False
+        ctx.is_causal = is_causal
+        ctx.scale_factor = scale_factor
+
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.has_mask:
+            query, key, value, output, logsumexp, attn_mask = ctx.saved_tensors
+        else:
+            query, key, value, output, logsumexp = ctx.saved_tensors
+            attn_mask = None
+
+        # Compute gradients
+        dQ, dK, dV = _C.backward(
+            grad_output, query, key, value, output, logsumexp, ctx.is_causal, attn_mask
+        )
+
+        # If we scaled the query in forward, scale the gradient back
+        if ctx.scale_factor != 1.0:
+            dQ = dQ * ctx.scale_factor
+
+        # Return gradients (None for is_causal, scale, and attn_mask since they're not tensors or don't need grad)
+        return dQ, dK, dV, None, None, None
+
+
+def flash_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    is_causal: bool = False,
+    scale: Optional[float] = None,
+    attn_mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Compute scaled dot-product attention using Flash Attention on MPS.
+
+    This function provides O(N) memory complexity instead of O(N²) by using
+    tiled computation, allowing much longer sequences on limited GPU memory.
+
+    Supports both forward and backward passes for training.
+
+    Args:
+        query: Query tensor of shape (B, num_heads, seq_len, head_dim)
+        key: Key tensor of shape (B, num_heads, seq_len, head_dim)
+        value: Value tensor of shape (B, num_heads, seq_len, head_dim)
+        is_causal: If True, applies causal masking (for autoregressive models)
+        scale: Scaling factor for attention scores. Default: 1/sqrt(head_dim)
+        attn_mask: Optional boolean attention mask of shape (B, 1, seq_len_q, seq_len_kv)
+                   or (B, num_heads, seq_len_q, seq_len_kv). True values indicate
+                   positions to be masked (not attended to).
+
+    Returns:
+        Output tensor of shape (B, num_heads, seq_len, head_dim)
+
+    Example:
+        >>> import torch
+        >>> from mps_flash_attn import flash_attention
+        >>> q = torch.randn(2, 8, 4096, 64, device='mps', dtype=torch.float16)
+        >>> k = torch.randn(2, 8, 4096, 64, device='mps', dtype=torch.float16)
+        >>> v = torch.randn(2, 8, 4096, 64, device='mps', dtype=torch.float16)
+        >>> out = flash_attention(q, k, v)
+
+        # With gradients:
+        >>> q.requires_grad = True
+        >>> out = flash_attention(q, k, v)
+        >>> out.sum().backward()  # Computes dQ
+
+        # With attention mask:
+        >>> mask = torch.zeros(2, 1, 4096, 4096, dtype=torch.bool, device='mps')
+        >>> mask[:, :, :, 2048:] = True  # mask out second half of keys
+        >>> out = flash_attention(q, k, v, attn_mask=mask)
+    """
+    if not _HAS_MFA:
+        raise RuntimeError(
+            f"MPS Flash Attention C++ extension not available: {_IMPORT_ERROR}\n"
+            "Please rebuild with: pip install -e ."
+        )
+
+    if not torch.backends.mps.is_available():
+        raise RuntimeError("MPS not available")
+
+    # Validate device
+    if query.device.type != 'mps':
+        raise ValueError("query must be on MPS device")
+    if key.device.type != 'mps':
+        raise ValueError("key must be on MPS device")
+    if value.device.type != 'mps':
+        raise ValueError("value must be on MPS device")
+    if attn_mask is not None and attn_mask.device.type != 'mps':
+        raise ValueError("attn_mask must be on MPS device")
+
+    # Fast path: inference mode (no grad) - skip autograd overhead and don't save tensors
+    if not torch.is_grad_enabled() or (not query.requires_grad and not key.requires_grad and not value.requires_grad):
+        # Apply scale if provided
+        if scale is not None:
+            default_scale = 1.0 / math.sqrt(query.shape[-1])
+            if abs(scale - default_scale) > 1e-6:
+                scale_factor = scale / default_scale
+                query = query * scale_factor
+
+        # Forward only - no logsumexp needed, no tensors saved
+        return _C.forward(query, key, value, is_causal, attn_mask)
+
+    # Use autograd function for gradient support
+    return FlashAttentionFunction.apply(query, key, value, is_causal, scale, attn_mask)
+
+
+def replace_sdpa():
+    """
+    Monkey-patch torch.nn.functional.scaled_dot_product_attention to use
+    Flash Attention on MPS devices.
+
+    Call this at the start of your script to automatically use Flash Attention
+    for all attention operations.
+    """
+    import torch.nn.functional as F
+
+    original_sdpa = F.scaled_dot_product_attention
+
+    def patched_sdpa(query, key, value, attn_mask=None, dropout_p=0.0,
+                     is_causal=False, scale=None):
+        # Use MFA for MPS tensors without dropout
+        # Only use MFA for seq_len >= 1024 where it outperforms PyTorch's math backend
+        # For shorter sequences, PyTorch's simpler matmul+softmax approach is faster
+        if (query.device.type == 'mps' and
+            dropout_p == 0.0 and
+            _HAS_MFA and
+            query.shape[2] >= 1024):
+            try:
+                # Convert float mask to bool mask if needed
+                # PyTorch SDPA uses additive masks (0 = attend, -inf = mask)
+                # MFA uses boolean masks (False/0 = attend, True/non-zero = mask)
+                mfa_mask = None
+                if attn_mask is not None:
+                    if attn_mask.dtype == torch.bool:
+                        # Boolean mask: True means masked (don't attend)
+                        mfa_mask = attn_mask
+                    else:
+                        # Float mask: typically -inf for masked positions, 0 for unmasked
+                        # Convert: positions with large negative values -> True (masked)
+                        mfa_mask = attn_mask < -1e4
+                return flash_attention(query, key, value, is_causal=is_causal, scale=scale, attn_mask=mfa_mask)
+            except Exception:
+                # Fall back to original on any error
+                pass
+
+        return original_sdpa(query, key, value, attn_mask, dropout_p, is_causal, scale=scale)
+
+    F.scaled_dot_product_attention = patched_sdpa
+    print("MPS Flash Attention: Patched F.scaled_dot_product_attention")
+
+
+def precompile():
+    """
+    Pre-compile Metal kernels for common configurations.
+
+    Call this once after installation to eliminate runtime compilation overhead.
+    Pre-compiled kernels are cached to disk and loaded instantly on subsequent runs.
+
+    This compiles kernels for:
+    - Sequence lengths: 64, 128, 256, 512, 1024, 2048, 4096, 8192
+    - Head dimensions: 32, 48, 64, 80, 96, 128
+    - Both fp32 and fp16 precision
+
+    Total: 96 kernel configurations
+    """
+    if not _HAS_MFA:
+        raise RuntimeError(f"MPS Flash Attention not available: {_IMPORT_ERROR}")
+
+    import ctypes
+    import os
+
+    # Load the Swift bridge directly
+    bridge_path = os.environ.get("MFA_BRIDGE_PATH")
+    if not bridge_path:
+        # Try common locations
+        module_dir = os.path.dirname(__file__)
+        candidates = [
+            os.path.join(module_dir, "lib", "libMFABridge.dylib"),  # Bundled in wheel
+            os.path.join(module_dir, "..", "swift-bridge", ".build", "release", "libMFABridge.dylib"),
+            os.path.join(module_dir, "libMFABridge.dylib"),
+        ]
+        for path in candidates:
+            if os.path.exists(path):
+                bridge_path = path
+                break
+
+    if not bridge_path or not os.path.exists(bridge_path):
+        raise RuntimeError("Cannot find libMFABridge.dylib. Set MFA_BRIDGE_PATH environment variable.")
+
+    lib = ctypes.CDLL(bridge_path)
+    lib.mfa_precompile()
+    print("\nPre-compilation complete! Kernels cached to disk.")
+
+
+def clear_cache():
+    """Clear the pre-compiled kernel cache."""
+    if not _HAS_MFA:
+        raise RuntimeError(f"MPS Flash Attention not available: {_IMPORT_ERROR}")
+
+    import ctypes
+    import os
+
+    bridge_path = os.environ.get("MFA_BRIDGE_PATH")
+    if bridge_path and os.path.exists(bridge_path):
+        lib = ctypes.CDLL(bridge_path)
+        lib.mfa_clear_cache()
+        print("Cache cleared.")

mps_flash_attn/csrc/mps_flash_attn.mm

@@ -0,0 +1,569 @@
+/**
+ * MPS Flash Attention - PyTorch C++ Extension
+ *
+ * Bridges PyTorch MPS tensors to the MFA Swift library.
+ */
+
+#include <torch/extension.h>
+#include <ATen/mps/MPSStream.h>
+#include <ATen/native/mps/OperationUtils.h>
+
+#import <Metal/Metal.h>
+#import <Foundation/Foundation.h>
+
+#include <dlfcn.h>
+#include <string>
+#include <vector>
+
+// ============================================================================
+// MFA Bridge Function Types
+// ============================================================================
+
+typedef bool (*mfa_init_fn)();
+typedef void* (*mfa_create_kernel_fn)(int32_t, int32_t, int32_t, bool, bool, bool, bool);
+// New zero-sync encode functions that take PyTorch's command encoder
+// Added mask_ptr and mask_offset parameters
+typedef bool (*mfa_forward_encode_fn)(void*, void*, void*, void*, void*, void*, void*, void*,
+                                       int64_t, int64_t, int64_t, int64_t, int64_t, int64_t,
+                                       int32_t, int32_t);
+typedef bool (*mfa_backward_encode_fn)(void*, void*, void*, void*, void*, void*, void*, void*, void*, void*, void*, void*, void*,
+                                        int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t,
+                                        int32_t, int32_t);
+// Legacy sync functions (fallback)
+typedef bool (*mfa_forward_fn)(void*, void*, void*, void*, void*, void*, void*,
+                                int64_t, int64_t, int64_t, int64_t, int64_t, int64_t,
+                                int32_t, int32_t);
+typedef bool (*mfa_backward_fn)(void*, void*, void*, void*, void*, void*, void*, void*, void*, void*, void*, void*,
+                                 int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t,
+                                 int32_t, int32_t);
+typedef void (*mfa_release_kernel_fn)(void*);
+
+// Global function pointers
+static mfa_init_fn g_mfa_init = nullptr;
+static mfa_create_kernel_fn g_mfa_create_kernel = nullptr;
+static mfa_forward_encode_fn g_mfa_forward_encode = nullptr;
+static mfa_backward_encode_fn g_mfa_backward_encode = nullptr;
+static mfa_forward_fn g_mfa_forward = nullptr;
+static mfa_backward_fn g_mfa_backward = nullptr;
+static mfa_release_kernel_fn g_mfa_release_kernel = nullptr;
+static void* g_dylib_handle = nullptr;
+static bool g_initialized = false;
+
+// ============================================================================
+// Load MFA Bridge Library
+// ============================================================================
+
+// Get the directory containing this shared library
+static std::string get_module_dir() {
+    Dl_info info;
+    if (dladdr((void*)get_module_dir, &info) && info.dli_fname) {
+        std::string path(info.dli_fname);
+        size_t last_slash = path.rfind('/');
+        if (last_slash != std::string::npos) {
+            return path.substr(0, last_slash);
+        }
+    }
+    return ".";
+}
+
+static bool load_mfa_bridge() {
+    if (g_dylib_handle) return true;
+
+    // First check environment variable (highest priority)
+    const char* mfa_path = getenv("MFA_BRIDGE_PATH");
+    if (mfa_path) {
+        g_dylib_handle = dlopen(mfa_path, RTLD_NOW);
+        if (g_dylib_handle) return true;
+    }
+
+    // Get the directory containing this extension module
+    std::string module_dir = get_module_dir();
+
+    // Try paths relative to the module directory
+    std::vector<std::string> paths = {
+        module_dir + "/lib/libMFABridge.dylib",  // Bundled in wheel
+        module_dir + "/../swift-bridge/.build/release/libMFABridge.dylib",  // Dev build
+        "libMFABridge.dylib",  // Current directory fallback
+    };
+
+    for (const auto& path : paths) {
+        g_dylib_handle = dlopen(path.c_str(), RTLD_NOW);
+        if (g_dylib_handle) break;
+    }
+
+    if (!g_dylib_handle) {
+        throw std::runtime_error(
+            "Failed to load libMFABridge.dylib. Set MFA_BRIDGE_PATH environment variable.");
+    }
+
+    // Load function pointers
+    g_mfa_init = (mfa_init_fn)dlsym(g_dylib_handle, "mfa_init");
+    g_mfa_create_kernel = (mfa_create_kernel_fn)dlsym(g_dylib_handle, "mfa_create_kernel");
+    g_mfa_forward_encode = (mfa_forward_encode_fn)dlsym(g_dylib_handle, "mfa_forward_encode");
+    g_mfa_backward_encode = (mfa_backward_encode_fn)dlsym(g_dylib_handle, "mfa_backward_encode");
+    g_mfa_forward = (mfa_forward_fn)dlsym(g_dylib_handle, "mfa_forward");
+    g_mfa_backward = (mfa_backward_fn)dlsym(g_dylib_handle, "mfa_backward");
+    g_mfa_release_kernel = (mfa_release_kernel_fn)dlsym(g_dylib_handle, "mfa_release_kernel");
+
+    // Require at least init, create_kernel, forward_encode (for zero-sync path)
+    if (!g_mfa_init || !g_mfa_create_kernel || !g_mfa_forward_encode) {
+        throw std::runtime_error("Failed to load MFA bridge functions");
+    }
+
+    return true;
+}
+
+// ============================================================================
+// Get MTLBuffer from PyTorch MPS Tensor
+// ============================================================================
+
+struct BufferInfo {
+    id<MTLBuffer> buffer;
+    int64_t byte_offset;
+};
+
+static BufferInfo getBufferInfo(const at::Tensor& tensor) {
+    TORCH_CHECK(tensor.device().is_mps(), "Tensor must be on MPS device");
+    TORCH_CHECK(tensor.is_contiguous(), "Tensor must be contiguous");
+
+    // Get the underlying Metal buffer (covers entire storage)
+    id<MTLBuffer> buffer = at::native::mps::getMTLBufferStorage(tensor);
+
+    // Calculate byte offset: storage_offset() is in elements, multiply by element size
+    int64_t element_size = tensor.element_size();
+    int64_t byte_offset = tensor.storage_offset() * element_size;
+
+    return {buffer, byte_offset};
+}
+
+// ============================================================================
+// Kernel Cache
+// ============================================================================
+
+struct KernelCacheKey {
+    int64_t seq_len_q;
+    int64_t seq_len_kv;
+    int64_t head_dim;
+    bool low_precision;
+    bool low_precision_outputs;
+    bool causal;
+    bool has_mask;
+
+    bool operator==(const KernelCacheKey& other) const {
+        return seq_len_q == other.seq_len_q &&
+               seq_len_kv == other.seq_len_kv &&
+               head_dim == other.head_dim &&
+               low_precision == other.low_precision &&
+               low_precision_outputs == other.low_precision_outputs &&
+               causal == other.causal &&
+               has_mask == other.has_mask;
+    }
+};
+
+struct KernelCacheKeyHash {
+    size_t operator()(const KernelCacheKey& k) const {
+        return std::hash<int64_t>()(k.seq_len_q) ^
+               (std::hash<int64_t>()(k.seq_len_kv) << 1) ^
+               (std::hash<int64_t>()(k.head_dim) << 2) ^
+               (std::hash<bool>()(k.low_precision) << 3) ^
+               (std::hash<bool>()(k.low_precision_outputs) << 4) ^
+               (std::hash<bool>()(k.causal) << 5) ^
+               (std::hash<bool>()(k.has_mask) << 6);
+    }
+};
+
+static std::unordered_map<KernelCacheKey, void*, KernelCacheKeyHash> g_kernel_cache;
+
+static void* get_or_create_kernel(int64_t seq_q, int64_t seq_kv, int64_t head_dim, bool low_prec, bool low_prec_outputs, bool causal, bool has_mask) {
+    KernelCacheKey key{seq_q, seq_kv, head_dim, low_prec, low_prec_outputs, causal, has_mask};
+
+    auto it = g_kernel_cache.find(key);
+    if (it != g_kernel_cache.end()) {
+        return it->second;
+    }
+
+    void* kernel = g_mfa_create_kernel(
+        static_cast<int32_t>(seq_q),
+        static_cast<int32_t>(seq_kv),
+        static_cast<int32_t>(head_dim),
+        low_prec,
+        low_prec_outputs,
+        causal,
+        has_mask
+    );
+
+    if (!kernel) {
+        throw std::runtime_error("Failed to create MFA kernel");
+    }
+
+    g_kernel_cache[key] = kernel;
+    return kernel;
+}
+
+// ============================================================================
+// Flash Attention Forward
+// ============================================================================
+
+std::tuple<at::Tensor, at::Tensor> mps_flash_attention_forward_with_lse(
+    const at::Tensor& query,   // (B, H, N, D)
+    const at::Tensor& key,     // (B, H, N, D)
+    const at::Tensor& value,   // (B, H, N, D)
+    bool is_causal,
+    const c10::optional<at::Tensor>& attn_mask  // Optional (B, 1, N_q, N_kv) or (B, H, N_q, N_kv)
+) {
+    // Initialize MFA on first call
+    if (!g_initialized) {
+        load_mfa_bridge();
+        if (!g_mfa_init()) {
+            throw std::runtime_error("Failed to initialize MFA");
+        }
+        g_initialized = true;
+    }
+
+    // Validate inputs
+    TORCH_CHECK(query.dim() == 4, "Query must be 4D (B, H, N, D)");
+    TORCH_CHECK(key.dim() == 4, "Key must be 4D (B, H, N, D)");
+    TORCH_CHECK(value.dim() == 4, "Value must be 4D (B, H, N, D)");
+    TORCH_CHECK(query.device().is_mps(), "Query must be on MPS device");
+    TORCH_CHECK(key.device().is_mps(), "Key must be on MPS device");
+    TORCH_CHECK(value.device().is_mps(), "Value must be on MPS device");
+
+    const int64_t batch_size = query.size(0);
+    const int64_t num_heads_q = query.size(1);
+    const int64_t num_heads_kv = key.size(1);
+    const int64_t seq_len_q = query.size(2);
+    const int64_t head_dim = query.size(3);
+    const int64_t seq_len_kv = key.size(2);
+
+    TORCH_CHECK(key.size(0) == batch_size && value.size(0) == batch_size,
+                "Batch size mismatch");
+    TORCH_CHECK(key.size(3) == head_dim && value.size(3) == head_dim,
+                "Head dimension mismatch");
+    TORCH_CHECK(key.size(1) == value.size(1),
+                "K and V must have same number of heads");
+
+    // Handle GQA (Grouped Query Attention): expand K/V if fewer heads than Q
+    const int64_t num_heads = num_heads_q;
+    at::Tensor k_expanded, v_expanded;
+
+    if (num_heads_kv != num_heads_q) {
+        // GQA: num_heads_q must be divisible by num_heads_kv
+        TORCH_CHECK(num_heads_q % num_heads_kv == 0,
+                    "num_heads_q (", num_heads_q, ") must be divisible by num_heads_kv (", num_heads_kv, ")");
+        int64_t repeat_factor = num_heads_q / num_heads_kv;
+
+        // Expand K and V to match Q's head count: (B, H_kv, S, D) -> (B, H_q, S, D)
+        // Use repeat_interleave for proper GQA expansion
+        k_expanded = key.repeat_interleave(repeat_factor, /*dim=*/1);
+        v_expanded = value.repeat_interleave(repeat_factor, /*dim=*/1);
+    } else {
+        k_expanded = key;
+        v_expanded = value;
+    }
+
+    // Determine precision
+    bool low_precision = (query.scalar_type() == at::kHalf ||
+                          query.scalar_type() == at::kBFloat16);
+
+    // For fp16 inputs, we can now output directly to fp16 (no extra conversion needed!)
+    bool low_precision_outputs = low_precision;
+
+    // Make inputs contiguous
+    auto q = query.contiguous();
+    auto k = k_expanded.contiguous();
+    auto v = v_expanded.contiguous();
+
+    // Handle attention mask
+    bool has_mask = attn_mask.has_value();
+    at::Tensor mask;
+    if (has_mask) {
+        mask = attn_mask.value();
+        TORCH_CHECK(mask.dim() == 4, "Attention mask must be 4D (B, H or 1, N_q, N_kv)");
+        TORCH_CHECK(mask.device().is_mps(), "Attention mask must be on MPS device");
+        // Convert to bool/uint8 if needed - kernel expects uchar (0 = attend, non-0 = mask out)
+        if (mask.scalar_type() == at::kBool) {
+            // Convert bool to uint8 for Metal compatibility
+            mask = mask.to(at::kByte);
+        }
+        TORCH_CHECK(mask.scalar_type() == at::kByte,
+                    "Attention mask must be bool or uint8");
+        // Expand mask heads if needed (B, 1, N_q, N_kv) -> (B, H, N_q, N_kv)
+        if (mask.size(1) == 1 && num_heads > 1) {
+            mask = mask.expand({batch_size, num_heads, seq_len_q, seq_len_kv});
+        }
+        mask = mask.contiguous();
+    }
+
+    // Allocate output in the appropriate precision
+    // With lowPrecisionOutputs=true, MFA writes FP16 directly
+    at::Tensor output;
+    if (low_precision_outputs) {
+        output = at::empty({batch_size, num_heads, seq_len_q, head_dim},
+                           query.options().dtype(at::kHalf));
+    } else {
+        output = at::empty({batch_size, num_heads, seq_len_q, head_dim},
+                           query.options().dtype(at::kFloat));
+    }
+
+    // Allocate logsumexp (for backward pass, always fp32)
+    auto logsumexp = at::empty({batch_size, num_heads, seq_len_q},
+                                query.options().dtype(at::kFloat));
+
+    // Get or create kernel with matching output precision and causal mode
+    void* kernel = get_or_create_kernel(seq_len_q, seq_len_kv, head_dim, low_precision, low_precision_outputs, is_causal, has_mask);
+
+    // Get Metal buffers with byte offsets
+    auto q_info = getBufferInfo(q);
+    auto k_info = getBufferInfo(k);
+    auto v_info = getBufferInfo(v);
+    auto o_info = getBufferInfo(output);
+    auto l_info = getBufferInfo(logsumexp);
+
+    // Mask buffer info (may be nullptr if no mask)
+    BufferInfo mask_info = {nil, 0};
+    if (has_mask) {
+        mask_info = getBufferInfo(mask);
+    }
+
+    // Use PyTorch's MPS stream command encoder for zero-sync integration
+    @autoreleasepool {
+        auto stream = at::mps::getCurrentMPSStream();
+
+        // Get PyTorch's shared command encoder - this is the key for zero-sync!
+        // All our dispatches go onto the same encoder that PyTorch uses.
+        id<MTLComputeCommandEncoder> encoder = stream->commandEncoder();
+
+        // Execute MFA using the shared encoder (no sync needed!)
+        bool success = g_mfa_forward_encode(
+            kernel,
+            (__bridge void*)encoder,  // PyTorch's shared command encoder
+            (__bridge void*)q_info.buffer,
+            (__bridge void*)k_info.buffer,
+            (__bridge void*)v_info.buffer,
+            (__bridge void*)o_info.buffer,
+            (__bridge void*)l_info.buffer,
+            has_mask ? (__bridge void*)mask_info.buffer : nullptr,
+            q_info.byte_offset,
+            k_info.byte_offset,
+            v_info.byte_offset,
+            o_info.byte_offset,
+            l_info.byte_offset,
+            mask_info.byte_offset,
+            static_cast<int32_t>(batch_size),
+            static_cast<int32_t>(num_heads)
+        );
+
+        if (!success) {
+            throw std::runtime_error("MFA forward pass failed");
+        }
+
+        // No commit needed - PyTorch will commit when it needs the results
+        // The encoder stays open for coalescing more kernels
+    }
+
+    // Output is already in the correct dtype (fp16 or fp32)
+    // Return both output and logsumexp (needed for backward pass)
+    return std::make_tuple(output, logsumexp);
+}
+
+// Simple forward that only returns output (for inference)
+at::Tensor mps_flash_attention_forward(
+    const at::Tensor& query,
+    const at::Tensor& key,
+    const at::Tensor& value,
+    bool is_causal,
+    const c10::optional<at::Tensor>& attn_mask
+) {
+    auto [output, logsumexp] = mps_flash_attention_forward_with_lse(query, key, value, is_causal, attn_mask);
+    return output;
+}
+
+// ============================================================================
+// Flash Attention Backward
+// ============================================================================
+
+std::tuple<at::Tensor, at::Tensor, at::Tensor> mps_flash_attention_backward(
+    const at::Tensor& grad_output,  // (B, H, N, D)
+    const at::Tensor& query,        // (B, H, N, D)
+    const at::Tensor& key,          // (B, H, N, D)
+    const at::Tensor& value,        // (B, H, N, D)
+    const at::Tensor& output,       // (B, H, N, D)
+    const at::Tensor& logsumexp,    // (B, H, N)
+    bool is_causal,
+    const c10::optional<at::Tensor>& attn_mask  // Optional (B, 1, N_q, N_kv) or (B, H, N_q, N_kv)
+) {
+    // Initialize MFA on first call
+    if (!g_initialized) {
+        load_mfa_bridge();
+        if (!g_mfa_init()) {
+            throw std::runtime_error("Failed to initialize MFA");
+        }
+        g_initialized = true;
+    }
+
+    // Validate inputs
+    TORCH_CHECK(grad_output.dim() == 4, "grad_output must be 4D (B, H, N, D)");
+    TORCH_CHECK(query.dim() == 4, "Query must be 4D (B, H, N, D)");
+    TORCH_CHECK(key.dim() == 4, "Key must be 4D (B, H, N, D)");
+    TORCH_CHECK(value.dim() == 4, "Value must be 4D (B, H, N, D)");
+    TORCH_CHECK(output.dim() == 4, "Output must be 4D (B, H, N, D)");
+    TORCH_CHECK(logsumexp.dim() == 3, "Logsumexp must be 3D (B, H, N)");
+
+    const int64_t batch_size = query.size(0);
+    const int64_t num_heads = query.size(1);
+    const int64_t seq_len_q = query.size(2);
+    const int64_t head_dim = query.size(3);
+    const int64_t seq_len_kv = key.size(2);
+
+    // Determine precision
+    bool low_precision = (query.scalar_type() == at::kHalf ||
+                          query.scalar_type() == at::kBFloat16);
+    bool low_precision_outputs = low_precision;
+
+    // Handle attention mask
+    bool has_mask = attn_mask.has_value();
+    at::Tensor mask;
+    if (has_mask) {
+        mask = attn_mask.value();
+        TORCH_CHECK(mask.dim() == 4, "Attention mask must be 4D (B, H or 1, N_q, N_kv)");
+        TORCH_CHECK(mask.device().is_mps(), "Attention mask must be on MPS device");
+        if (mask.scalar_type() == at::kBool) {
+            mask = mask.to(at::kByte);
+        }
+        TORCH_CHECK(mask.scalar_type() == at::kByte,
+                    "Attention mask must be bool or uint8");
+        if (mask.size(1) == 1 && num_heads > 1) {
+            mask = mask.expand({batch_size, num_heads, seq_len_q, seq_len_kv});
+        }
+        mask = mask.contiguous();
+    }
+
+    // Make inputs contiguous and upcast to fp32 for numerical stability
+    // The backward pass accumulates many small values, so fp32 precision is critical
+    auto q = query.contiguous().to(at::kFloat);
+    auto k = key.contiguous().to(at::kFloat);
+    auto v = value.contiguous().to(at::kFloat);
+    auto o = output.contiguous().to(at::kFloat);
+    auto dO = grad_output.contiguous().to(at::kFloat);
+    auto lse = logsumexp.contiguous();
+
+    // Get or create kernel - always use fp32 for backward pass
+    void* kernel = get_or_create_kernel(seq_len_q, seq_len_kv, head_dim, false, false, is_causal, has_mask);
+
+    // Allocate D buffer (dO * O reduction, always fp32)
+    auto D = at::empty({batch_size, num_heads, seq_len_q},
+                        query.options().dtype(at::kFloat));
+
+    // Allocate gradients (always fp32 for numerical stability)
+    auto dQ = at::zeros({batch_size, num_heads, seq_len_q, head_dim},
+                         query.options().dtype(at::kFloat));
+    auto dK = at::zeros({batch_size, num_heads, seq_len_kv, head_dim},
+                         query.options().dtype(at::kFloat));
+    auto dV = at::zeros({batch_size, num_heads, seq_len_kv, head_dim},
+                         query.options().dtype(at::kFloat));
+
+    // Get Metal buffers with byte offsets
+    auto q_info = getBufferInfo(q);
+    auto k_info = getBufferInfo(k);
+    auto v_info = getBufferInfo(v);
+    auto o_info = getBufferInfo(o);
+    auto do_info = getBufferInfo(dO);
+    auto l_info = getBufferInfo(lse);
+    auto d_info = getBufferInfo(D);
+    auto dq_info = getBufferInfo(dQ);
+    auto dk_info = getBufferInfo(dK);
+    auto dv_info = getBufferInfo(dV);
+
+    // Mask buffer info (may be nullptr if no mask)
+    BufferInfo mask_info = {nil, 0};
+    if (has_mask) {
+        mask_info = getBufferInfo(mask);
+    }
+
+    // Use PyTorch's MPS stream command encoder for zero-sync integration
+    @autoreleasepool {
+        auto stream = at::mps::getCurrentMPSStream();
+
+        // Get PyTorch's shared command encoder
+        id<MTLComputeCommandEncoder> encoder = stream->commandEncoder();
+
+        bool success = g_mfa_backward_encode(
+            kernel,
+            (__bridge void*)encoder,  // PyTorch's shared command encoder
+            (__bridge void*)q_info.buffer,
+            (__bridge void*)k_info.buffer,
+            (__bridge void*)v_info.buffer,
+            (__bridge void*)o_info.buffer,
+            (__bridge void*)do_info.buffer,
+            (__bridge void*)l_info.buffer,
+            (__bridge void*)d_info.buffer,
+            (__bridge void*)dq_info.buffer,
+            (__bridge void*)dk_info.buffer,
+            (__bridge void*)dv_info.buffer,
+            has_mask ? (__bridge void*)mask_info.buffer : nullptr,
+            q_info.byte_offset,
+            k_info.byte_offset,
+            v_info.byte_offset,
+            o_info.byte_offset,
+            do_info.byte_offset,
+            l_info.byte_offset,
+            d_info.byte_offset,
+            dq_info.byte_offset,
+            dk_info.byte_offset,
+            dv_info.byte_offset,
+            mask_info.byte_offset,
+            static_cast<int32_t>(batch_size),
+            static_cast<int32_t>(num_heads)
+        );
+
+        if (!success) {
+            throw std::runtime_error("MFA backward pass failed");
+        }
+
+        // No commit needed - PyTorch will commit when it needs the results
+    }
+
+    // Convert gradients back to input dtype if needed
+    if (low_precision) {
+        dQ = dQ.to(query.scalar_type());
+        dK = dK.to(query.scalar_type());
+        dV = dV.to(query.scalar_type());
+    }
+
+    return std::make_tuple(dQ, dK, dV);
+}
+
+// ============================================================================
+// Python Bindings
+// ============================================================================
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.doc() = "MPS Flash Attention - Metal accelerated attention for Apple Silicon";
+
+    m.def("forward", &mps_flash_attention_forward,
+          "Flash Attention forward pass (returns output only)",
+          py::arg("query"),
+          py::arg("key"),
+          py::arg("value"),
+          py::arg("is_causal") = false,
+          py::arg("attn_mask") = py::none());
+
+    m.def("forward_with_lse", &mps_flash_attention_forward_with_lse,
+          "Flash Attention forward pass (returns output and logsumexp for backward)",
+          py::arg("query"),
+          py::arg("key"),
+          py::arg("value"),
+          py::arg("is_causal") = false,
+          py::arg("attn_mask") = py::none());
+
+    m.def("backward", &mps_flash_attention_backward,
+          "Flash Attention backward pass",
+          py::arg("grad_output"),
+          py::arg("query"),
+          py::arg("key"),
+          py::arg("value"),
+          py::arg("output"),
+          py::arg("logsumexp"),
+          py::arg("is_causal") = false,
+          py::arg("attn_mask") = py::none());
+}

mps_flash_attn/kernels/manifest.json

@@ -0,0 +1,27 @@
+{
+  "version": "1.0",
+  "files": [
+    "06c421e7a01418cf64aafa07f6b1df0558148583959c596d9a7ce260987f89f0.metallib",
+    "adc0f77ff05156bbda8fe78afd9ba8a8d3c890ba8fea0902ae79a6ae8c4f04c3.metallib",
+    "eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f_4096_4096.bin",
+    "eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f_2048_2048.bin",
+    "eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f_8192_8192.bin",
+    "eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f_1024_1024.bin",
+    "ac4573fb201e92310867c59bd569a8ae68f859d60a9352d9d4d5d41c1547c83c.metallib",
+    "73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2_512_512.bin",
+    "a5e2c5c401e3872af0899c1fb3e30b5f52a6070fc49c9dac02982cc1c2f25849.metallib",
+    "73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2.metallib",
+    "975aece2b4d3d78035be08a0735a7deacf2e544adee5af81c9c0a3a42a926129.metallib",
+    "2ca9312d1151f792e1a95617db9186928300e3d0ffbe016f0ad53b62ab840bac.metallib",
+    "73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2_2048_2048.bin",
+    "09b9615289be632fdf05444004a0b3b67fb1b70b05a7e0fce8e0ba3a95e3921c.metallib",
+    "73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2_4096_4096.bin",
+    "0c36461301fb52cbad786d0642b020ad2bfc7229b487ccb5dff44d198423b347.metallib",
+    "73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2_1024_1024.bin",
+    "73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2_8192_8192.bin",
+    "771935bf47d248650e287da82bc82e04bff7c4c52964823e7a12462ccd23408e.metallib",
+    "eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f_512_512.bin",
+    "eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f.metallib",
+    "f08fe0efd72e055177e068154dae01e08c4d52d3cb883330a04f1431d274aece.metallib"
+  ]
+}

mps_flash_attn-0.1.7.dist-info/METADATA

@@ -0,0 +1,270 @@
+Metadata-Version: 2.4
+Name: mps-flash-attn
+Version: 0.1.7
+Summary: Flash Attention for PyTorch on Apple Silicon (M1/M2/M3/M4)
+Author: imperatormk
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/mpsops/mps-flash-attention
+Project-URL: Repository, https://github.com/mpsops/mps-flash-attention
+Project-URL: Issues, https://github.com/mpsops/mps-flash-attention/issues
+Keywords: flash-attention,apple-silicon,pytorch,mps,metal,transformer,attention
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0.0
+Dynamic: license-file
+
+# MPS Flash Attention
+
+Flash Attention for PyTorch on Apple Silicon (M1/M2/M3/M4).
+
+**O(N) memory** instead of O(N²), enabling 8K+ sequence lengths on unified memory.
+
+## Features
+
+- **Forward pass**: 2-5x faster than PyTorch SDPA
+- **Backward pass**: Full gradient support for training (fp32 precision)
+- **Causal masking**: Native kernel support (only 5% overhead)
+- **Attention masks**: Full boolean mask support for arbitrary masking patterns
+- **FP16/FP32**: Native fp16 output (no conversion overhead)
+- **Pre-compiled kernels**: Zero-compilation cold start (~6ms)
+
+## Performance
+
+Tested on M1 Max, N=2048, B=4, H=8, D=64:
+
+| Operation | MPS Flash Attn | PyTorch SDPA | Speedup |
+|-----------|----------------|--------------|---------|
+| Forward | 5.3ms | 15ms | 2.8x |
+| Forward+Backward | 55ms | 108ms | 2.0x |
+| Memory | 80MB | 592MB | 7.4x less |
+
+## Installation
+
+### Prerequisites
+
+- macOS 14+ (Sonoma) or macOS 15+ (Sequoia)
+- Xcode Command Line Tools (`xcode-select --install`)
+- Python 3.10+ with PyTorch 2.0+
+
+### Build from source
+
+```bash
+# Clone with submodules
+git clone --recursive https://github.com/mpsops/mps-flash-attention.git
+cd mps-flash-attention
+
+# Build Swift bridge
+cd swift-bridge
+swift build -c release
+cd ..
+
+# Install Python package
+pip install -e .
+```
+
+### Set environment variable
+
+```bash
+export MFA_BRIDGE_PATH=/path/to/mps-flash-attention/swift-bridge/.build/release/libMFABridge.dylib
+```
+
+## Usage
+
+### Basic usage
+
+```python
+from mps_flash_attn import flash_attention
+
+# Standard attention (B, H, N, D)
+q = torch.randn(2, 8, 4096, 64, device='mps', dtype=torch.float16)
+k = torch.randn(2, 8, 4096, 64, device='mps', dtype=torch.float16)
+v = torch.randn(2, 8, 4096, 64, device='mps', dtype=torch.float16)
+
+out = flash_attention(q, k, v)
+```
+
+### Causal masking (for autoregressive models)
+
+```python
+out = flash_attention(q, k, v, is_causal=True)
+```
+
+### Attention masks (for custom masking patterns)
+
+```python
+# Boolean mask: True = masked (don't attend), False = attend
+mask = torch.zeros(B, 1, N, N, dtype=torch.bool, device='mps')
+mask[:, :, :, 512:] = True  # Mask out positions after 512
+
+out = flash_attention(q, k, v, attn_mask=mask)
+```
+
+### Training with gradients
+
+```python
+q.requires_grad = True
+k.requires_grad = True
+v.requires_grad = True
+
+out = flash_attention(q, k, v, is_causal=True)
+loss = out.sum()
+loss.backward()  # Computes dQ, dK, dV
+```
+
+### Drop-in replacement for SDPA
+
+```python
+from mps_flash_attn import replace_sdpa
+
+# Monkey-patch F.scaled_dot_product_attention
+replace_sdpa()
+
+# Now all attention ops use Flash Attention on MPS
+```
+
+## Architecture
+
+```
++----------------------------------------------------------+
+|                    Python API                            |
+|              mps_flash_attn/__init__.py                  |
+|         (flash_attention, autograd Function)             |
++----------------------------+-----------------------------+
+                             |
++----------------------------v-----------------------------+
+|                 C++ Extension                            |
+|            mps_flash_attn/csrc/mps_flash_attn.mm         |
+|    (PyTorch bindings, MTLBuffer handling, offsets)       |
++----------------------------+-----------------------------+
+                             | dlopen + dlsym
++----------------------------v-----------------------------+
+|                 Swift Bridge                             |
+|         swift-bridge/Sources/MFABridge/                  |
+|   (MFABridge.swift, MetallibCache.swift)                 |
+|   @_cdecl exports: mfa_init, mfa_create_kernel,          |
+|                    mfa_forward, mfa_backward             |
++----------------------------+-----------------------------+
+                             |
++----------------------------v-----------------------------+
+|              Metal Flash Attention                       |
+|    metal-flash-attention/Sources/FlashAttention/         |
+|     (AttentionDescriptor, AttentionKernel, etc.)         |
+|                                                          |
+|   Generates Metal shader source at runtime,              |
+|   compiles to .metallib, caches pipelines                |
++----------------------------------------------------------+
+```
+
+## Project Structure
+
+```
+mps-flash-attention/
+├── mps_flash_attn/              # Python package
+│   ├── __init__.py              # Public API (flash_attention, replace_sdpa)
+│   ├── csrc/
+│   │   └── mps_flash_attn.mm    # PyTorch C++ extension
+│   └── kernels/                 # Pre-compiled metallibs (optional)
+│
+├── swift-bridge/                # Swift -> C bridge
+│   ├── Package.swift
+│   └── Sources/MFABridge/
+│       ├── MFABridge.swift      # C-callable API (@_cdecl)
+│       └── MetallibCache.swift  # Disk caching for metallibs
+│
+├── metal-flash-attention/       # Upstream (git submodule)
+│   └── Sources/FlashAttention/
+│       └── Attention/
+│           ├── AttentionDescriptor/  # Problem configuration
+│           ├── AttentionKernel/      # Metal shader generation
+│           └── ...
+│
+├── scripts/
+│   └── build_metallibs.py       # Pre-compile kernels for distribution
+│
+└── setup.py                     # Python package setup
+```
+
+## Changes from upstream metal-flash-attention
+
+We made the following modifications to `metal-flash-attention`:
+
+### 1. macOS 15+ compatibility (MTLLibraryCompiler.swift)
+
+Apple restricted `__asm` in runtime-compiled Metal shaders on macOS 15. We added a fallback that uses `xcrun metal` CLI compilation when runtime compilation fails.
+
+### 2. Causal masking support
+
+Added `causal` flag to AttentionDescriptor and kernel generation:
+
+- `AttentionDescriptor.swift`: Added `causal: Bool` property
+- `AttentionKernelDescriptor.swift`: Added `causal: Bool` property
+- `AttentionKernel.swift`: Added `causal` field
+- `AttentionKernel+Softmax.swift`: Added `maskCausal()` function
+- `AttentionKernel+Source.swift`: Added causal masking to forward/backward loops
+
+## Next Steps
+
+### 1. PR to upstream metal-flash-attention
+
+The macOS 15 fix and causal masking should be contributed back:
+
+```bash
+cd metal-flash-attention
+git checkout -b macos15-causal-support
+# Commit changes to:
+#   - Sources/FlashAttention/Utilities/MTLLibraryCompiler.swift (new file)
+#   - Sources/FlashAttention/Attention/AttentionDescriptor/*.swift
+#   - Sources/FlashAttention/Attention/AttentionKernel/*.swift
+git push origin macos15-causal-support
+# Open PR at https://github.com/philipturner/metal-flash-attention
+```
+
+### 2. Publish mps-flash-attention to PyPI
+
+```bash
+# Add pyproject.toml with proper metadata
+# Build wheel with pre-compiled Swift bridge
+python -m build
+twine upload dist/*
+```
+
+### 3. Pre-compile kernels for zero cold start
+
+```bash
+python scripts/build_metallibs.py
+# Copies metallibs to mps_flash_attn/kernels/
+# These get shipped with the wheel
+```
+
+## Current Status (Jan 2025)
+
+**Working:**
+- Forward pass (fp16/fp32)
+- Backward pass (dQ, dK, dV gradients)
+- Causal masking
+- Metallib disk caching
+- Pipeline binary caching (MTLBinaryArchive)
+
+**Known limitations:**
+- Sequence length must be divisible by block size (typically 64)
+- Head dimension: Best with 32, 64, 96, 128
+- No dropout
+
+## Credits
+
+- [metal-flash-attention](https://github.com/philipturner/metal-flash-attention) by Philip Turner
+- [Flash Attention](https://arxiv.org/abs/2205.14135) paper by Tri Dao et al.
+
+## License
+
+MIT

mps_flash_attn-0.1.7.dist-info/RECORD

@@ -0,0 +1,31 @@
+mps_flash_attn/_C.cpython-314-darwin.so,sha256=Wy33hq4OLXRvTI_HrBdnml7_MVqhIAC6LfHl2501WcY,266984
+mps_flash_attn/__init__.py,sha256=HCz_Qs65RcxGD3FZ6tirBfxvFo0kk4BiOhFO-3QV_Pg,9980
+mps_flash_attn/csrc/mps_flash_attn.mm,sha256=CSqtWw15FbQ3O50SgoOmL2ZsG2rKnAu8s8rE0MzxWas,22567
+mps_flash_attn/kernels/06c421e7a01418cf64aafa07f6b1df0558148583959c596d9a7ce260987f89f0.metallib,sha256=_oig6f2I6ZxBCKWbJF3ofmZMySm8gB399_M-lD2NOfM,13747
+mps_flash_attn/kernels/09b9615289be632fdf05444004a0b3b67fb1b70b05a7e0fce8e0ba3a95e3921c.metallib,sha256=1fsmVvB5EubhN-y6s5CB-eVk_wuO2tfrabiQTwXvJJc,13171
+mps_flash_attn/kernels/0c36461301fb52cbad786d0642b020ad2bfc7229b487ccb5dff44d198423b347.metallib,sha256=5WKo_yAU-PgmulBUQhnzvt0DZRteVmo4-nc4U-T6G2g,17507
+mps_flash_attn/kernels/2ca9312d1151f792e1a95617db9186928300e3d0ffbe016f0ad53b62ab840bac.metallib,sha256=OehFzjntFdIgEIvE1EW2sGxDEzUgreATvI-fm5-HjaI,12723
+mps_flash_attn/kernels/73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2.metallib,sha256=FUcHAIglSvbw6aavubMKzBH6PO5Z1TkUhKO4IqIZiLQ,14083
+mps_flash_attn/kernels/73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2_1024_1024.bin,sha256=4qmQhKBCDqZEydpRaPhuF0WBCiDXfALSW8SbAE9HUps,36496
+mps_flash_attn/kernels/73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2_2048_2048.bin,sha256=NB1ex3aP-SARN_u4I24laoe_H_rbfRBLyDgDnZH4IYY,36496
+mps_flash_attn/kernels/73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2_4096_4096.bin,sha256=uZpt_PpesdfmS7ywNrGqFZecNy4q0yXKQIz0rKu-9yg,36496
+mps_flash_attn/kernels/73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2_512_512.bin,sha256=dj6zpVaJtduxjWXg3J8bTuZo_LfAEutNGc0quZXQqSA,36496
+mps_flash_attn/kernels/73254c55475c6b7f7b009f095398994b1f9ae8215beafcf810f100357ccc99b2_8192_8192.bin,sha256=gwGlY6Od0Ub1u3JfRLDkksYFQR-aT3A-igFDYaBDEYA,36496
+mps_flash_attn/kernels/771935bf47d248650e287da82bc82e04bff7c4c52964823e7a12462ccd23408e.metallib,sha256=bfG4lBjhtTXAJEInsCBHfE9CzwUkLAZWduLy-jr6alA,12819
+mps_flash_attn/kernels/975aece2b4d3d78035be08a0735a7deacf2e544adee5af81c9c0a3a42a926129.metallib,sha256=rPg7TwGLETfD5-GJHAHy084Wza0N5BcsifgxKEQ-HkU,12707
+mps_flash_attn/kernels/a5e2c5c401e3872af0899c1fb3e30b5f52a6070fc49c9dac02982cc1c2f25849.metallib,sha256=2q-7oJOeS8GR66iJJisbfNbIVuY44E41wwh_UgfBlug,17203
+mps_flash_attn/kernels/ac4573fb201e92310867c59bd569a8ae68f859d60a9352d9d4d5d41c1547c83c.metallib,sha256=QVQMAdhrinNhTJ6vYHgeWWz9-tv4Kp3Z9O4tJvLfjLM,13747
+mps_flash_attn/kernels/adc0f77ff05156bbda8fe78afd9ba8a8d3c890ba8fea0902ae79a6ae8c4f04c3.metallib,sha256=v20uzcsN3HJAxopiFfu90cKJ_zNeUb13OZzm1ubGXaA,13747
+mps_flash_attn/kernels/eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f.metallib,sha256=YhZNagEBf68RxgO8ZqTujfu0YjcUD85dq0Iv5fC8QLE,13171
+mps_flash_attn/kernels/eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f_1024_1024.bin,sha256=13GAdfAiz-3wWcKhvvMbydtoo9pCMXqUsYujQSv5A_8,34112
+mps_flash_attn/kernels/eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f_2048_2048.bin,sha256=DOicBen4ow2FSk2Jv2KPiDZ5-6H93-LvwafvOPqfoF0,34112
+mps_flash_attn/kernels/eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f_4096_4096.bin,sha256=CyRMosbr9mYqxARS3ta_v2cBoYuqnHmEwwQ6aaXsDUY,34112
+mps_flash_attn/kernels/eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f_512_512.bin,sha256=IHdTXUtoGnKKYBXeO-Yj5quXy91lgmfxhrD4Bn_yGzA,34112
+mps_flash_attn/kernels/eab4f40de4b0ebd2765b41c25dba7ccab5db4abf6a6eb87d76fff7b5e0829b2f_8192_8192.bin,sha256=4BjRprnMycnhZql9829R6FS3HW30jejuDJM9p9vzVPs,34112
+mps_flash_attn/kernels/f08fe0efd72e055177e068154dae01e08c4d52d3cb883330a04f1431d274aece.metallib,sha256=qyOaQtRVwL_Wc6GGdu6z-ftf0iX84XexuY09-lNLl5o,13747
+mps_flash_attn/kernels/manifest.json,sha256=d5MkE_BjqDQuMNm1jZiwWkQKfB-yfFml3lLSeR-wCLo,1867
+mps_flash_attn-0.1.7.dist-info/licenses/LICENSE,sha256=F_XmXSab2O-hHcqLpYJWeFaqB6GA_qiTEN23p2VfZWU,1237
+mps_flash_attn-0.1.7.dist-info/METADATA,sha256=eIKQbJo__DR8opFhCTok4KsjNZyJynRXaL_r_riA7M4,8634
+mps_flash_attn-0.1.7.dist-info/WHEEL,sha256=uAzMRtb2noxPlbYLbRgeD25pPgKOo3k59IS71Dg5Qjs,110
+mps_flash_attn-0.1.7.dist-info/top_level.txt,sha256=zbArDcWhJDnJfMUKnOUhs5TjsMgSxa2GzOlscTRfobE,15
+mps_flash_attn-0.1.7.dist-info/RECORD,,

mps_flash_attn-0.1.7.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.10.2)
+Root-Is-Purelib: false
+Tag: cp314-cp314-macosx_15_0_arm64
+

mps_flash_attn-0.1.7.dist-info/licenses/LICENSE

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

mps_flash_attn-0.1.7.dist-info/top_level.txt

@@ -0,0 +1 @@
+mps_flash_attn