torch-ane 0.1.0__tar.gz

torch_ane-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Ivan Shivalov
+
+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.

torch_ane-0.1.0/PKG-INFO

@@ -0,0 +1,167 @@
+Metadata-Version: 2.4
+Name: torch-ane
+Version: 0.1.0
+Summary: Run PyTorch models on the Apple Neural Engine via Core AI: compile any nn.Module, LLMs with KV cache, Stable Diffusion
+Author: Ivan Shivalov
+License-Expression: MIT
+Project-URL: Repository, https://github.com/intexcor/torch-ane
+Project-URL: Issues, https://github.com/intexcor/torch-ane/issues
+Keywords: pytorch,apple-silicon,neural-engine,ane,core-ai,coreml,llm,stable-diffusion
+Classifier: Development Status :: 3 - Alpha
+Classifier: Operating System :: MacOS
+Classifier: Programming Language :: Python :: 3
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.10
+Requires-Dist: numpy
+Requires-Dist: coreai-torch>=0.4.3
+Requires-Dist: coreai-opt>=0.3
+Provides-Extra: llm
+Requires-Dist: transformers>=5; extra == "llm"
+Requires-Dist: gguf; extra == "llm"
+Provides-Extra: diffusion
+Requires-Dist: diffusers>=0.40; extra == "diffusion"
+Requires-Dist: transformers>=5; extra == "diffusion"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pyflakes; extra == "dev"
+Dynamic: license-file
+
+# torch-ane
+
+Run PyTorch models on the **Apple Neural Engine** through Apple's Core AI (macOS 27+).
+
+Pass in a plain `nn.Module`, or a model from transformers or diffusers. torch-ane rewrites it
+into a form the ANE runs fast and correctly: fp16, static shapes, and the layout and ops from
+Apple's ANE guidance. It also works around the ANE's numeric quirks. You get back a module
+that behaves like the original.
+
+```python
+import torch_ane
+
+m = torch_ane.compile(model.eval(), example_inputs=(x,))  # any nn.Module -> ANE
+y = m(x)                                                   # torch tensors in, torch tensors out
+```
+
+On an M4 MacBook, ANE prefill and convolutions are 2–3× faster than the GPU (MLX/MPS).
+Decoding is limited by memory bandwidth, and the ANE reads ~57 GB/s.
+
+## Install
+
+Requires macOS 27+ on Apple silicon, and Python 3.11+.
+
+```sh
+pip install torch-ane              # core: compile()
+pip install "torch-ane[llm]"       # + CausalLM / ChunkedCausalLM (transformers, gguf)
+pip install "torch-ane[diffusion]" # + StableDiffusion (diffusers)
+```
+
+Core AI's Python packages (`coreai-core`, `coreai-torch`, `coreai-opt`) come from PyPI
+as dependencies.
+
+## What's inside
+
+```python
+import torch_ane
+
+torch_ane.available()          # is Core AI usable here
+
+# Any module. The first call per input shape builds a program (seconds); then it's cached on disk.
+m = torch_ane.compile(model, example_inputs=(x,), device="ane")   # or device="gpu"
+
+# A whole LLM on the ANE (up to ~1-3B). The KV cache is Core AI state, so it is never copied.
+lm = torch_ane.CausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", batch=True)
+ids = lm.generate(prompt_ids, max_new_tokens=64, eos=tok.eos_token_id)
+outs = lm.generate_batch([ids1, ids2, ids3], max_new_tokens=64)   # up to 8 streams per step
+logits = lm.feed(more_ids)                                        # append to the cache -> last logits
+
+# A large model from GGUF, in 4-layer chunks with 4-bit palettes (Devstral 24B etc.)
+big = torch_ane.ChunkedCausalLM.from_gguf("model.gguf", "config_dir/")
+
+# Stable Diffusion: text encoder, UNet and VAE on the ANE
+sd = torch_ane.StableDiffusion("segmind/tiny-sd")
+image, timings = sd("a cute robot", steps=20, seed=0)
+```
+
+| | |
+|---|---|
+| `compile(model, example_inputs=None, device="ane", validate=True, compress=None)` | → `CompiledModule`. `validate` checks the first call of each shape against eager. If they disagree, that shape falls back to GPU, then to eager. `compress="int8"` stores Linear/Conv2d weights in int8 |
+| `CompiledModule` | torch module over Core AI programs, one program per input shape (up to 8) |
+| `CausalLM(hf_model, max_len=2048, compress=None, batch=False, ...)` | prefill in chunks of 256, decode in windows of 8, a ladder of context sizes 512…`max_len` in one program |
+| `ChunkedCausalLM.from_gguf(path, config)` | 4-bit palettes, plus int8 where the GGUF is more precise (Q6_K). First build takes ~35 min for 24B |
+| `StableDiffusion(model_id, device="ane", size=512)` | diffusers pipeline with all three parts on the ANE |
+| `passes` | the building blocks: `export_for_ane`, `swap_linear`, `split_wide`, `chunk_mlp`, `palettize_lut`, `quantize_int8_modules`, `calibrate_rescale`, `use_composite_ops`, … |
+| `CACHE_DIR`, `prune_cache()`, `clear_system_cache()` | build cache at `~/.cache/torch_ane` (`TORCH_ANE_CACHE`). It is capped at `TORCH_ANE_CACHE_GB` (40) and least recently used builds are removed. `clear_system_cache()` wipes Core AI's own compile cache, which can only be cleared as a whole |
+
+## What it does to your model
+
+These passes need no knowledge of the specific model:
+
+1. RMSNorm, found by a numerical check rather than by name, becomes a Core AI composite op.
+2. Everything that doesn't depend on the input (masks, RoPE tables) is precomputed in fp32.
+3. A causal mask becomes `is_causal`.
+4. SiLU becomes `x·0.5·(1 + tanh(x/2))`. The native ANE SiLU is off by ~5% near zero, which
+   wrecks GroupNorm→SiLU in diffusion UNets.
+5. The model is converted to fp16, the only float type the ANE runs.
+6. MLPs whose projection weight exceeds 4 MB are split along the hidden dimension. Each chunk
+   of `up` feeds its own chunk of `down`, and the results are summed. Otherwise the weight plus
+   the `[S, hidden]` activation doesn't fit the ANE's on-chip memory (~32 MB on M4), and the
+   ANE stalls on the bus at ~50 GB/s. This cuts the Whisper large encoder by 24% and
+   Qwen2.5-0.5B prefill by 14%. Disable it with `rules=()`.
+
+LLMs get two more passes:
+- Per-projection 2^k scales, calibrated on real text, so small activations don't become fp16
+  subnormals while large ones don't overflow.
+- int16 positions, embeddings on the CPU, and logits for one row only.
+
+UNet/VAE attention uses Apple's layout: BC1S, heads split, `einsum` with no transposes.
+
+## Speed (MacBook M4, 24 GB)
+
+| task | ANE | for comparison |
+|---|---|---|
+| ResNet-18, batch 16 | 4.5 ms | MLX 24.5, MPS 32.7 |
+| BERT-base 8×128 | 20 ms | MLX 71, MPS 75 |
+| Whisper large-v3-turbo encoder | 285 ms (int8: 264) | |
+| Qwen2.5-0.5B prefill, 512 tokens | 9065 tok/s | mlx-lm 3279 |
+| Qwen2.5-0.5B decode, fp16 | 54 tok/s; 8 streams: 351 | mlx-lm bf16 70 |
+| Devstral Small 24B decode | 4.8 tok/s | mlx-lm 4-bit 6.7 |
+| tiny-sd 512×512, 20 steps | 4.9 s | VAE 3× faster than GPU |
+
+## Examples
+
+- [`examples/compile_bert.py`](examples/compile_bert.py): BERT-base, eager vs ANE
+- [`examples/generate.py`](examples/generate.py): chat with Qwen2.5 on the ANE
+- [`examples/chat_gguf.py`](examples/chat_gguf.py): chat with a 24B GGUF model, with the engine in a restartable process
+- [`examples/imagine.py`](examples/imagine.py): Stable Diffusion
+
+## Limitations
+
+- Shapes are static, so a new input shape means a new build.
+- The ANE computes in fp16 only (weights can be palettized or int8). `-inf` in masks is
+  replaced by −40000.
+- int8 weights in Conv1d make the ANE compiler fail silently, and the whole function moves to
+  the GPU. `compress="int8"` therefore quantizes only Linear/Conv2d.
+- The coreai-core Python wrapper aborts the process after 16 384 allocated buffers
+  ([apple/coreai-torch#75](https://github.com/apple/coreai-torch/issues/75)).
+  `ChunkedCausalLM.steps_left()` lets you restart ahead of time; see `examples/chat_gguf.py`.
+- The ANE unloads programs after ~10 s idle, which makes the next call ~3× slower. Use
+  `ChunkedCausalLM.keep_warm()` to avoid this.
+- Copying an input into Core AI runs at ~0.5 GB/s, about 10 ms per 4 MB input. The public API
+  has no way around this.
+- Code comments are in Russian.
+
+## Tests
+
+```sh
+pip install -e ".[llm,diffusion,dev]"
+pytest
+```
+
+## License
+
+MIT

torch_ane-0.1.0/README.md

@@ -0,0 +1,134 @@
+# torch-ane
+
+Run PyTorch models on the **Apple Neural Engine** through Apple's Core AI (macOS 27+).
+
+Pass in a plain `nn.Module`, or a model from transformers or diffusers. torch-ane rewrites it
+into a form the ANE runs fast and correctly: fp16, static shapes, and the layout and ops from
+Apple's ANE guidance. It also works around the ANE's numeric quirks. You get back a module
+that behaves like the original.
+
+```python
+import torch_ane
+
+m = torch_ane.compile(model.eval(), example_inputs=(x,))  # any nn.Module -> ANE
+y = m(x)                                                   # torch tensors in, torch tensors out
+```
+
+On an M4 MacBook, ANE prefill and convolutions are 2–3× faster than the GPU (MLX/MPS).
+Decoding is limited by memory bandwidth, and the ANE reads ~57 GB/s.
+
+## Install
+
+Requires macOS 27+ on Apple silicon, and Python 3.11+.
+
+```sh
+pip install torch-ane              # core: compile()
+pip install "torch-ane[llm]"       # + CausalLM / ChunkedCausalLM (transformers, gguf)
+pip install "torch-ane[diffusion]" # + StableDiffusion (diffusers)
+```
+
+Core AI's Python packages (`coreai-core`, `coreai-torch`, `coreai-opt`) come from PyPI
+as dependencies.
+
+## What's inside
+
+```python
+import torch_ane
+
+torch_ane.available()          # is Core AI usable here
+
+# Any module. The first call per input shape builds a program (seconds); then it's cached on disk.
+m = torch_ane.compile(model, example_inputs=(x,), device="ane")   # or device="gpu"
+
+# A whole LLM on the ANE (up to ~1-3B). The KV cache is Core AI state, so it is never copied.
+lm = torch_ane.CausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", batch=True)
+ids = lm.generate(prompt_ids, max_new_tokens=64, eos=tok.eos_token_id)
+outs = lm.generate_batch([ids1, ids2, ids3], max_new_tokens=64)   # up to 8 streams per step
+logits = lm.feed(more_ids)                                        # append to the cache -> last logits
+
+# A large model from GGUF, in 4-layer chunks with 4-bit palettes (Devstral 24B etc.)
+big = torch_ane.ChunkedCausalLM.from_gguf("model.gguf", "config_dir/")
+
+# Stable Diffusion: text encoder, UNet and VAE on the ANE
+sd = torch_ane.StableDiffusion("segmind/tiny-sd")
+image, timings = sd("a cute robot", steps=20, seed=0)
+```
+
+| | |
+|---|---|
+| `compile(model, example_inputs=None, device="ane", validate=True, compress=None)` | → `CompiledModule`. `validate` checks the first call of each shape against eager. If they disagree, that shape falls back to GPU, then to eager. `compress="int8"` stores Linear/Conv2d weights in int8 |
+| `CompiledModule` | torch module over Core AI programs, one program per input shape (up to 8) |
+| `CausalLM(hf_model, max_len=2048, compress=None, batch=False, ...)` | prefill in chunks of 256, decode in windows of 8, a ladder of context sizes 512…`max_len` in one program |
+| `ChunkedCausalLM.from_gguf(path, config)` | 4-bit palettes, plus int8 where the GGUF is more precise (Q6_K). First build takes ~35 min for 24B |
+| `StableDiffusion(model_id, device="ane", size=512)` | diffusers pipeline with all three parts on the ANE |
+| `passes` | the building blocks: `export_for_ane`, `swap_linear`, `split_wide`, `chunk_mlp`, `palettize_lut`, `quantize_int8_modules`, `calibrate_rescale`, `use_composite_ops`, … |
+| `CACHE_DIR`, `prune_cache()`, `clear_system_cache()` | build cache at `~/.cache/torch_ane` (`TORCH_ANE_CACHE`). It is capped at `TORCH_ANE_CACHE_GB` (40) and least recently used builds are removed. `clear_system_cache()` wipes Core AI's own compile cache, which can only be cleared as a whole |
+
+## What it does to your model
+
+These passes need no knowledge of the specific model:
+
+1. RMSNorm, found by a numerical check rather than by name, becomes a Core AI composite op.
+2. Everything that doesn't depend on the input (masks, RoPE tables) is precomputed in fp32.
+3. A causal mask becomes `is_causal`.
+4. SiLU becomes `x·0.5·(1 + tanh(x/2))`. The native ANE SiLU is off by ~5% near zero, which
+   wrecks GroupNorm→SiLU in diffusion UNets.
+5. The model is converted to fp16, the only float type the ANE runs.
+6. MLPs whose projection weight exceeds 4 MB are split along the hidden dimension. Each chunk
+   of `up` feeds its own chunk of `down`, and the results are summed. Otherwise the weight plus
+   the `[S, hidden]` activation doesn't fit the ANE's on-chip memory (~32 MB on M4), and the
+   ANE stalls on the bus at ~50 GB/s. This cuts the Whisper large encoder by 24% and
+   Qwen2.5-0.5B prefill by 14%. Disable it with `rules=()`.
+
+LLMs get two more passes:
+- Per-projection 2^k scales, calibrated on real text, so small activations don't become fp16
+  subnormals while large ones don't overflow.
+- int16 positions, embeddings on the CPU, and logits for one row only.
+
+UNet/VAE attention uses Apple's layout: BC1S, heads split, `einsum` with no transposes.
+
+## Speed (MacBook M4, 24 GB)
+
+| task | ANE | for comparison |
+|---|---|---|
+| ResNet-18, batch 16 | 4.5 ms | MLX 24.5, MPS 32.7 |
+| BERT-base 8×128 | 20 ms | MLX 71, MPS 75 |
+| Whisper large-v3-turbo encoder | 285 ms (int8: 264) | |
+| Qwen2.5-0.5B prefill, 512 tokens | 9065 tok/s | mlx-lm 3279 |
+| Qwen2.5-0.5B decode, fp16 | 54 tok/s; 8 streams: 351 | mlx-lm bf16 70 |
+| Devstral Small 24B decode | 4.8 tok/s | mlx-lm 4-bit 6.7 |
+| tiny-sd 512×512, 20 steps | 4.9 s | VAE 3× faster than GPU |
+
+## Examples
+
+- [`examples/compile_bert.py`](examples/compile_bert.py): BERT-base, eager vs ANE
+- [`examples/generate.py`](examples/generate.py): chat with Qwen2.5 on the ANE
+- [`examples/chat_gguf.py`](examples/chat_gguf.py): chat with a 24B GGUF model, with the engine in a restartable process
+- [`examples/imagine.py`](examples/imagine.py): Stable Diffusion
+
+## Limitations
+
+- Shapes are static, so a new input shape means a new build.
+- The ANE computes in fp16 only (weights can be palettized or int8). `-inf` in masks is
+  replaced by −40000.
+- int8 weights in Conv1d make the ANE compiler fail silently, and the whole function moves to
+  the GPU. `compress="int8"` therefore quantizes only Linear/Conv2d.
+- The coreai-core Python wrapper aborts the process after 16 384 allocated buffers
+  ([apple/coreai-torch#75](https://github.com/apple/coreai-torch/issues/75)).
+  `ChunkedCausalLM.steps_left()` lets you restart ahead of time; see `examples/chat_gguf.py`.
+- The ANE unloads programs after ~10 s idle, which makes the next call ~3× slower. Use
+  `ChunkedCausalLM.keep_warm()` to avoid this.
+- Copying an input into Core AI runs at ~0.5 GB/s, about 10 ms per 4 MB input. The public API
+  has no way around this.
+- Code comments are in Russian.
+
+## Tests
+
+```sh
+pip install -e ".[llm,diffusion,dev]"
+pytest
+```
+
+## License
+
+MIT

torch_ane-0.1.0/pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+requires = ["setuptools>=69"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torch-ane"
+version = "0.1.0"
+description = "Run PyTorch models on the Apple Neural Engine via Core AI: compile any nn.Module, LLMs with KV cache, Stable Diffusion"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [{ name = "Ivan Shivalov" }]
+requires-python = ">=3.11"
+dependencies = ["torch>=2.10", "numpy", "coreai-torch>=0.4.3", "coreai-opt>=0.3"]
+keywords = ["pytorch", "apple-silicon", "neural-engine", "ane", "core-ai", "coreml", "llm", "stable-diffusion"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Operating System :: MacOS",
+    "Programming Language :: Python :: 3",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+[project.optional-dependencies]
+llm = ["transformers>=5", "gguf"]
+diffusion = ["diffusers>=0.40", "transformers>=5"]
+dev = ["pytest", "pyflakes"]
+
+[project.urls]
+Repository = "https://github.com/intexcor/torch-ane"
+Issues = "https://github.com/intexcor/torch-ane/issues"
+
+[tool.setuptools.packages.find]
+include = ["torch_ane*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

torch_ane-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torch_ane-0.1.0/tests/test_torch_ane.py

@@ -0,0 +1,207 @@
+"""torch_ane: компиляция на Neural Engine, проходы, LLM с KV-кешем-состоянием."""
+import subprocess
+import sys
+
+import pytest
+import torch
+import torch.nn as nn
+
+import torch_ane
+
+needs_coreai = pytest.mark.skipif(not torch_ane.available(), reason="нет Core AI (macOS 27 + coreai-torch)")
+
+
+def test_no_mlx_dependency():
+    """Библиотека самостоятельна: импорт torch_ane и его LLM-части не тянет MLX и torch_mlx."""
+    code = ("import sys, torch_ane; torch_ane.CausalLM; torch_ane.passes.swap_linear; "
+            "print(any(m == 'mlx' or m.startswith(('mlx.', 'torch_mlx')) for m in sys.modules))")
+    out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True).stdout
+    assert out.strip() == "False"
+
+
+def test_prune_cache(tmp_path, monkeypatch):
+    """Кеш сборок не больше лимита: удаляются давно не использованные, keep и
+    недописанные (.partial) остаются."""
+    import os
+
+    from torch_ane import runtime
+    monkeypatch.setattr(runtime, "CACHE_DIR", tmp_path)
+    for i, name in enumerate(("a.aimodel", "b.aimodel", "c.aimodel", "d.partial.aimodel")):
+        d = tmp_path / name
+        d.mkdir()
+        (d / "w").write_bytes(b"x" * 1000)
+        os.utime(d, (1000 + i, 1000 + i))
+    freed = runtime.prune_cache(max_gb=2.5e-6, keep=(tmp_path / "a.aimodel",))  # лимит 2500 байт
+    assert freed == 2000 and sorted(p.name for p in tmp_path.iterdir()) == ["a.aimodel", "d.partial.aimodel"]
+
+
+def test_device_names():
+    m = nn.Linear(4, 4)
+    if torch_ane.available():
+        assert torch_ane.CompiledModule(m, device="ane").device == "neural_engine"
+        assert torch_ane.CompiledModule(m, device="gpu").device == "gpu"
+    with pytest.raises(ValueError):
+        torch_ane.runtime._device("tpu")
+
+
+@needs_coreai
+def test_compile_matches_eager():
+    torch.manual_seed(0)
+    m = nn.Sequential(nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 8)).eval()
+    x = torch.randn(4, 32)
+    cm = torch_ane.compile(m, example_inputs=(x,))
+    with torch.no_grad():
+        ref = m(x)
+    out = cm(x)
+    assert out.dtype == torch.float32 and ((out - ref).norm() / ref.norm()) < 2e-2
+
+
+@needs_coreai
+def test_compile_int8_with_conv1d():
+    """compress="int8" на ANE: int8 только в Linear/Conv2d. Сжатие готовой программы
+    квантовало и Conv1d — компилятор ANE падал, модель уходила на GPU с мусором на выходе."""
+    class Net(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv = nn.Conv1d(16, 64, 3, padding=1)
+            self.mlp = nn.Sequential(nn.Linear(64, 256), nn.GELU(), nn.Linear(256, 64))
+
+        def forward(self, x):  # [1, 16, T]
+            h = torch.nn.functional.gelu(self.conv(x)).transpose(1, 2)
+            return h + self.mlp(h)
+
+    torch.manual_seed(0)
+    m = Net().eval()
+    x = torch.randn(1, 16, 128)
+    cm = torch_ane.compile(m, example_inputs=(x,), validate=False, compress="int8")
+    with torch.no_grad():
+        ref = m(x)
+    out = cm(x)
+    if cm.ane_stats:
+        assert cm.ane_stats["int8"] == 2  # две Linear, Conv1d не тронута
+    assert ((out - ref).norm() / ref.norm()) < 3e-2
+
+
+@needs_coreai
+def test_ane_passes_llm():
+    """ANE-проходы на LLM из transformers без правок под модель: RMSNorm -> составная операция,
+    свёртка констант (маска, RoPE), причинная маска -> is_causal; следующий токен как у eager."""
+    from transformers import Qwen2Config, Qwen2ForCausalLM
+
+    class Prefill(nn.Module):
+        def __init__(self, lm):
+            super().__init__()
+            self.lm = lm
+
+        def forward(self, ids):
+            return self.lm.lm_head(self.lm.model(input_ids=ids, use_cache=False).last_hidden_state[:, -1])
+
+    torch.manual_seed(0)
+    lm = Qwen2ForCausalLM(Qwen2Config(hidden_size=128, intermediate_size=256, num_hidden_layers=2,
+                                      num_attention_heads=4, num_key_value_heads=2, vocab_size=1000)).eval()
+    model = Prefill(lm).eval()
+    ids = torch.randint(0, 1000, (1, 64))
+    m = torch_ane.compile(model, validate=False)
+    out = m(ids)
+    with torch.no_grad():
+        ref = model(ids)
+    if m.ane_stats:  # при загрузке из дискового кеша проходы не запускаются
+        assert {k: m.ane_stats[k] for k in ("rms_norm", "causal")} == {"rms_norm": 5, "causal": 2}
+        assert m.ane_stats["folded"] > 0
+    assert int(out.argmax()) == int(ref.argmax())
+    assert ((out.float() - ref).norm() / ref.norm()) < 5e-2
+
+
+def test_ane_rules_eager_equivalent():
+    """Правила Apple для ANE (Linear -> свёртка 1x1, внимание по головам) не меняют
+    математику модели: проверка в eager, без Core AI."""
+    from transformers import Qwen2Config, Qwen2ForCausalLM
+
+    from torch_ane import passes
+    torch.manual_seed(0)
+    lm = Qwen2ForCausalLM(Qwen2Config(hidden_size=64, intermediate_size=128, num_hidden_layers=2,
+                                      num_attention_heads=4, num_key_value_heads=2, vocab_size=100)).eval()
+    ids = torch.randint(0, 100, (1, 23))
+    with torch.no_grad():
+        ref = lm(ids).logits
+        for rules in (("conv",), ("conv_n111",), ("heads",), ("conv", "heads")):
+            m, st = passes.prepare(lm, rules)
+            assert st.get("conv", 0) or st.get("heads", 0)
+            torch.testing.assert_close(m(ids).logits, ref, rtol=1e-4, atol=1e-4)
+
+
+def test_chunk_mlp_eager_equivalent():
+    """MLP кусками скрытого измерения: Sequential (Whisper) и gate/up/down (Qwen) — та же
+    математика; маленькие MLP не трогаются."""
+    from transformers import Qwen2Config, Qwen2ForCausalLM
+
+    from torch_ane import passes
+    torch.manual_seed(0)
+    seq = nn.Sequential(nn.Sequential(nn.Linear(256, 1024), nn.GELU(), nn.Linear(1024, 256))).eval()
+    lm = Qwen2ForCausalLM(Qwen2Config(hidden_size=64, intermediate_size=256, num_hidden_layers=2,
+                                      num_attention_heads=4, num_key_value_heads=2, vocab_size=100)).eval()
+    x, ids = torch.randn(3, 256), torch.randint(0, 100, (1, 11))
+    with torch.no_grad():
+        ref_s, ref_l = seq(x), lm(ids).logits
+        m, st = passes.prepare(seq, ("mlp_chunk",))
+        assert st["mlp_chunk"] == 0  # вес 0.5 МБ — меньше порога
+        assert passes.chunk_mlp(m, max_bytes=128 << 10) == 1 and len(m[0].up) == 4
+        torch.testing.assert_close(m(x), ref_s, rtol=1e-4, atol=1e-5)
+        m, _ = passes.prepare(lm, ())
+        assert passes.chunk_mlp(m, max_bytes=8 << 10) == 2 and len(m.model.layers[0].mlp.up) == 4
+        torch.testing.assert_close(m(ids).logits, ref_l, rtol=1e-4, atol=1e-4)
+        m, _ = passes.prepare(lm, ())  # down_proj в обёртке масштаба (CausalLM)
+        passes.rescale_for_fp16(m, k=4)
+        assert passes.chunk_mlp(m, max_bytes=8 << 10) == 2 and m.model.layers[0].mlp.scale == 2 ** -4
+        torch.testing.assert_close(m(ids).logits, ref_l, rtol=1e-4, atol=1e-4)
+
+
+@needs_coreai
+def test_causal_lm_matches_eager():
+    """KV-кеш — состояния Core AI (схема Apple): prefill куском + декод окнами по 8
+    дают те же логиты, что transformers на всём контексте."""
+    from transformers import Qwen2Config, Qwen2ForCausalLM
+
+    import numpy as np
+
+    torch.manual_seed(0)
+    lm = Qwen2ForCausalLM(Qwen2Config(hidden_size=128, intermediate_size=256, num_hidden_layers=2,
+                                      num_attention_heads=4, num_key_value_heads=2, vocab_size=1000)).eval()
+    ids = torch.randint(0, 1000, (100,)).tolist()  # 37 + 63: переход со ступени 64 на 128
+    st = torch_ane.CausalLM(lm, max_len=256, prefill_q=32, ladder=(64, 128))
+    st.reset()
+    st._hist = []
+    got = [st.feed(ids[:37])] + [st.feed([t]) for t in ids[37:]]
+    with torch.no_grad():
+        ref = lm(torch.tensor([ids])).logits[0, 36:].numpy()
+    assert st.M == 128
+    for g, r in zip(got, ref):
+        assert np.linalg.norm(g - r) / np.linalg.norm(r) < 0.05
+
+
+@needs_coreai
+def test_causal_lm_batch_matches_single():
+    """Потоки в окне декода (generate_batch): у каждого свои позиции и маска, K/V вперемешку
+    в общем буфере — токены совпадают с одиночной генерацией каждого промпта."""
+    from transformers import Qwen2Config, Qwen2ForCausalLM
+
+    torch.manual_seed(0)
+    lm = Qwen2ForCausalLM(Qwen2Config(hidden_size=128, intermediate_size=256, num_hidden_layers=2,
+                                      num_attention_heads=4, num_key_value_heads=2, vocab_size=1000)).eval()
+    prompts = [torch.randint(0, 1000, (n,)).tolist() for n in (5, 40, 17)]  # 40 — два куска prefill
+    st = torch_ane.CausalLM(lm, max_len=256, prefill_q=32, ladder=(128,), batch=True)
+    assert st.generate_batch(prompts, 12) == [st.generate(p, 12) for p in prompts]
+
+
+@needs_coreai
+def test_silu_after_groupnorm():
+    """Родная SiLU на ANE врёт около нуля (~5% после GroupNorm); export_for_ane заменяет её
+    записью через tanh — связка GroupNorm -> SiLU -> свёртка совпадает с eager."""
+    torch.manual_seed(0)
+    m = nn.Sequential(nn.GroupNorm(8, 64), nn.SiLU(), nn.Conv2d(64, 64, 3, padding=1)).eval()
+    x = torch.randn(2, 64, 16, 16) * 3
+    with torch.no_grad():
+        ref = m(x)
+    cm = torch_ane.compile(m, validate=False)
+    out = cm(x)
+    assert ((out - ref).norm() / ref.norm()) < 5e-3

torch_ane-0.1.0/torch_ane/__init__.py

@@ -0,0 +1,34 @@
+"""torch_ane — PyTorch-модели на Apple Neural Engine через Core AI (macOS 27+).
+
+    import torch_ane
+
+    m = torch_ane.compile(model.eval(), example_inputs=(x,))     # любой nn.Module -> ANE
+    y = m(x)
+
+    lm = torch_ane.CausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
+    ids = lm.generate(prompt_ids, max_new_tokens=64)
+
+    sd = torch_ane.StableDiffusion("segmind/tiny-sd")
+    image, timings = sd("a cute robot", steps=20, seed=0)
+
+Нижний уровень — torch_ane.passes (проходы экспорта, раскладка Apple, палитры/int8,
+составные операции). Сборки кешируются в torch_ane.CACHE_DIR (~/.cache/torch_ane,
+переменная TORCH_ANE_CACHE). Библиотека не зависит от MLX.
+"""
+from . import passes
+from .runtime import (CACHE_DIR, CompiledModule, available, cleanup_mpsgraph_tmp, clear_system_cache, compile,
+                      compress_program, model_digest, prune_cache)
+
+__all__ = ["compile", "CompiledModule", "available", "compress_program", "model_digest", "cleanup_mpsgraph_tmp",
+           "CACHE_DIR", "prune_cache", "clear_system_cache", "passes", "CausalLM", "ChunkedCausalLM", "StableDiffusion"]
+
+
+def __getattr__(name):
+    # тяжёлые части (transformers, diffusers) — только по обращению
+    if name in ("CausalLM", "ChunkedCausalLM"):
+        from . import llm
+        return getattr(llm, name)
+    if name == "StableDiffusion":
+        from .diffusion import StableDiffusion
+        return StableDiffusion
+    raise AttributeError(f"module 'torch_ane' has no attribute {name!r}")