# mps-linear-attention

Metal kernels for Flash Linear Attention on Apple Silicon (MPS). Drop-in replacement for [`causal-conv1d`](https://github.com/Dao-AILab/causal-conv1d) and [`flash-linear-attention`](https://github.com/sustcsonglin/flash-linear-attention) — both CUDA-only — enabling the fast path for models like **Qwen3.5** on Mac.

## Install

```bash
pip install mps-linear-attn
```

Or from source:

```bash
git clone https://github.com/mpsops/mps-linear-attention
cd mps-linear-attention
pip install -e . --no-build-isolation
```

## Usage

Import before loading your model — that's it:

```python
import mps_linear_attn  # patches causal_conv1d + fla on import

from transformers import Qwen3_5ForConditionalGeneration, AutoTokenizer
import torch

model = Qwen3_5ForConditionalGeneration.from_pretrained(
    "Qwen/Qwen3.5-0.8B", dtype=torch.float16
).to("mps")
```

No code changes needed. The fast path activates automatically.

## What it does

Transformers checks for `causal_conv1d` and `flash-linear-attention` at import time to enable the fast path. Both require CUDA and fail silently on MPS, forcing a slow fallback.

This package:
1. Monkey-patches `sys.modules` with MPS-compatible implementations
2. Overrides `is_causal_conv1d_available()` and `is_flash_linear_attention_available()` to return `True`
3. Provides Metal GPU kernels for the hot paths

## Metal kernels

| Kernel | Description |
|--------|-------------|
| `causal_conv1d_fwd` | Depthwise causal 1D conv + optional SiLU, fp16/fp32 |
| `causal_conv1d_update` | Sliding-window state update + conv + SiLU (decode) |
| `delta_rule_recurrent_step` | Single-token gated delta rule state update |
| `delta_rule_recurrent_fused` | Fused T-step recurrent loop (prefill), single dispatch |

State is always `float32` for numerical precision. Q/K/V/beta are `float16`.

## Performance

Tested on Qwen3.5-0.8B, M1, macOS 26, PyTorch 2.7:

| | tok/s |
|---|---|
| Baseline (no fast path) | ~16 |
| mps-linear-attn | **~30** |

## Supported models

- **Qwen3.5** (all sizes) — uses GatedDeltaNet linear attention layers
- Any model using `fla.ops.gated_delta_rule` or `causal_conv1d`

## Requirements

- macOS (Apple Silicon — M1 or later)
- PyTorch >= 2.0 with MPS support
- Xcode Command Line Tools
