{"id":51319536,"url":"https://github.com/manishklach/mlx-metal-kernels","last_synced_at":"2026-07-01T11:31:32.225Z","repository":{"id":365999742,"uuid":"1274654754","full_name":"manishklach/mlx-metal-kernels","owner":"manishklach","description":"Experimental MLX custom Metal kernels for Apple Silicon — fast attention, decode, KV-cache, and future Mac GPU inference primitives.","archived":false,"fork":false,"pushed_at":"2026-06-19T19:33:41.000Z","size":76,"stargazers_count":1,"open_issues_count":12,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-06-19T21:16:07.276Z","etag":null,"topics":["apple-gpu","apple-silicon","attention","custom-kernels","deep-learning","flashattention","gpu-kernels","inference","kv-cache","llm","llm-inference","machine-learning","macos","metal","metal-kernels","mlx","mps","python","transformers"],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/manishklach.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":"docs/roadmap.md","authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-06-19T18:42:19.000Z","updated_at":"2026-06-19T19:33:44.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/manishklach/mlx-metal-kernels","commit_stats":null,"previous_names":["manishklach/mlx-metal-kernels"],"tags_count":null,"template":false,"template_full_name":null,"purl":"pkg:github/manishklach/mlx-metal-kernels","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/manishklach%2Fmlx-metal-kernels","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/manishklach%2Fmlx-metal-kernels/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/manishklach%2Fmlx-metal-kernels/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/manishklach%2Fmlx-metal-kernels/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/manishklach","download_url":"https://codeload.github.com/manishklach/mlx-metal-kernels/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/manishklach%2Fmlx-metal-kernels/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":35005409,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-07-01T02:00:05.325Z","response_time":130,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["apple-gpu","apple-silicon","attention","custom-kernels","deep-learning","flashattention","gpu-kernels","inference","kv-cache","llm","llm-inference","machine-learning","macos","metal","metal-kernels","mlx","mps","python","transformers"],"created_at":"2026-07-01T11:31:29.515Z","updated_at":"2026-07-01T11:31:32.216Z","avatar_url":"https://github.com/manishklach.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# MLX Metal Kernels\n\n[![CI](https://github.com/manishklach/mlx-metal-kernels/actions/workflows/ci.yml/badge.svg)](https://github.com/manishklach/mlx-metal-kernels/actions/workflows/ci.yml)\n![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue)\n[![License](https://img.shields.io/badge/license-Apache--2.0-blue)](LICENSE)\n\nExperimental MLX custom Metal kernels for Apple Silicon LLM inference.\n\nThis repository is a correctness-first kernel lab for exploring how far Apple Silicon GPUs can be pushed through MLX custom Metal kernels. It includes attention kernels, decode attention, KV-cache and paged KV-cache operations, q4/q8 quantized matvec kernels, RMSNorm, RoPE, SwiGLU, fused decode helpers, quantized decode blocks, quantized MLP blocks, experimental fused q4/q8 MLP kernels, toy transformer-layer benchmarks, autotuning infrastructure, and Llama-like model integration scaffolding.\n\nThe project is not a production inference engine. It is an experimental research repo focused on building, validating, benchmarking, and iterating on Mac-native transformer inference primitives.\n\n![Architecture](docs/architecture.svg)\n\nDesign principles:\n\n- correctness first\n- pure MLX reference path for every optimized backend\n- no performance claims without local Apple Silicon benchmark data\n- explicit experimental backend flags\n- shape-specialized kernels where useful\n- machine-specific autotuning instead of universal performance assumptions\n\n## What this repo explores\n\nThe repo currently covers several layers of the local LLM inference stack:\n\n- attention and FlashAttention-style streaming softmax\n- decode attention over contiguous and paged KV-cache layouts\n- threadgroup and simdgroup Metal experiments\n- q4/q8 dequantization and quantized matvec\n- multi-output tiled quantized GEMV\n- experimental fused q4/q8 MLP kernels\n- RMSNorm, RoPE, SwiGLU, and residual helpers\n- fused and quantized decode blocks\n- quantized MLP blocks\n- toy transformer-layer decode benchmarks\n- backend autotuning for Apple Silicon machines\n- Llama-like configuration and weight-layout scaffolding\n- GQA/MQA decode support\n- checkpoint-to-quantized packaging scaffold\n- tokenizer and sampling scaffold\n\n## What this project is not\n\nThis repo does not currently claim to outperform MLX native kernels, CUDA FlashAttention, or production inference engines.\n\nIt is not yet a full model runtime, tokenizer pipeline, checkpoint loader, or production Llama/Mistral inference engine.\n\nEnd-to-end generation from real model weights has not yet been benchmarked; current results are kernel- and block-level only.\n\nPerformance claims should only be made from benchmark data generated on a specific Apple Silicon machine.\n\n## Kernel Families\n\n### Attention\n\n- Reference MLX attention\n- Baseline custom Metal streaming attention\n- Row-parallel attention\n- Tiled K/V attention\n- Sparse sliding-window GQA/MQA attention\n- Threadgroup attention v2\n- Shape-specialized D=64/D=128 attention\n- Experimental `simdgroup_d64` attention\n\n### Decode and KV-cache\n\n- Contiguous KV-cache update\n- Decode attention\n- Paged KV-cache allocation and update\n- Paged decode attention\n- Decode loop helpers\n- Fused decode blocks from QKV\n- GQA/MQA decode-block composition\n- GQA/MQA prefill attention\n\n### Transformer primitives\n\n- RMSNorm\n- RoPE\n- SwiGLU\n- Residual add\n- RMSNorm + residual\n- QKV split\n- QKV split + RoPE\n- QKV + RoPE + cache update\n\n### Quantization\n\n- q4/q8 dequantization\n- q4/q8 decode matvec\n- parallel q4/q8 matvec\n- multi-output tiled q4/q8 matvec\n- quantized QKV projection\n- quantized output projection\n- quantized decode block\n- quantized MLP block\n- fused q4/q8 MLP experiments\n\n### Model-level scaffolding\n\n- Toy transformer-layer decode benchmark\n- Full Llama-like decode layer experiment\n- Multi-layer decode stack\n- Llama-like config\n- Weight-layout mapping helpers\n- Model-adapter scaffold\n- GQA/MQA utilities\n\n### Speculative decoding / MTP\n\n- SpeculativeConfig, DraftProposal, VerificationResult, SpeculativeGenerationResult\n- FixedDraftProposer, RandomDraftProposer, GreedySelfDraftProposer\n- PipelineTargetVerifier (sequential greedy verification)\n- SpeculativeGenerator (propose → verify → accept/reject → commit)\n- Accept/reject helpers: compute_accept_mask, verify_draft_tokens\n- Cache commit/rollback ops: commit_accepted_cache, discard_suffix\n- Synthetic MTP scaffold: MTPConfig, SyntheticMTPHead, mtp_propose_tokens\n- Pipeline integration via generate_speculative() (opt-in, existing generate() unchanged)\n\n### Parallel speculative verification (staged decode‑loop)\n\n- `ParallelVerificationConfig`, `ParallelVerificationPassResult`\n- `parallel_verify_tokens` — runs target model on all draft tokens in a staged decode loop, computes accept/reject mask, preserves committed cache isolation via a cloned staged KV‑cache\n- `commit_parallel_verification_cache` — finalises accepted positions into the committed cache\n- `ParallelTargetVerifier` — drop‑in target verifier wrapping the parallel path, usable as `target_verifier` on `SpeculativeGenerator`\n- `generate_speculative(verifier_mode=\"parallel\")` — pipeline entry point (mode=\"sequential\" preserves original behaviour)\n- Benchmark: `python benchmarks/bench_parallel_speculative_verify.py --verifier both` compares sequential vs. parallel verifier latency\n- Demo: `python examples/parallel_speculative_verify_demo.py`\n- **Current limitation**: verification runs as a staged decode loop rather than a true batched prefill, because continuation prefill (`start_position \u003e 0`) is not yet implemented. See `docs/parallel_speculative_verification.md` for details.\n\n### Async KV prefetch scheduler\n\nThe repo includes an experimental step-based KV prefetch scheduler for long-context offload experiments. It can queue KV block prefetch requests, simulate in-flight latency, prefetch sparse-attention blocks ahead of use, and report scheduler activity.\n\nCurrent scope:\n\n- prefetch request objects\n- deterministic step scheduler\n- in-flight/completed tracking\n- sparse decode lookahead planning\n- speculative draft lookahead planning\n- integration with long-context runtime reports\n- simulated latency benchmark\n\nOut of scope:\n\n- production async IO\n- DMA\n- real SSD streaming\n- automatic scheduling policy\n- kernel-level compute/IO overlap\n\nCommands:\n\n```\npython examples/kv_prefetch_scheduler_demo.py\npython benchmarks/bench_kv_prefetch_scheduler.py --seq-len 4096 --block-size 128 --num-layers 2 --num-kv-heads 8 --head-dim 128 --window-size 512 --sink-tokens 4 --lookahead-tokens 8 --max-in-flight 4 --simulated-latency-steps 2 --offload-ratio 0.75\n```\n\n### Flash/NAND KV offload tier scaffold\n\nThe repo includes an experimental KV-cache offload scaffold for long-context inference. It tracks hot/resident and cold/offloaded KV blocks, supports in-memory and local file-backed stores, and provides prefetch planning for sparse attention patterns.\n\nCurrent scope:\n\n- KV block metadata\n- residency maps\n- in-memory offload store\n- file-backed `.npy` offload store\n- offload/prefetch operations\n- sparse-attention residency guard\n- simulated benchmark\n\nOut of scope:\n\n- production SSD streaming\n- DMA\n- async IO scheduler\n- automatic offload policy in generation\n- model-quality claims\n\nCommands:\n```\npython examples/kv_offload_demo.py\npython examples/sparse_kv_offload_demo.py\npython benchmarks/bench_kv_offload_tier.py --seq-len 4096 --block-size 128 --num-layers 2 --num-kv-heads 8 --head-dim 128 --store memory --offload-ratio 0.75 --window-size 512 --sink-tokens 4\n```\n \n### Quantized KV-cache attention\n \nThe repo includes experimental q8/q4 KV-cache decode attention paths. K/V cache values can be stored in quantized form with per-token/head/group scales and dequantized inside GQA/MQA decode attention.\n \nCurrent scope:\n \n- q8 KV-cache quantization\n- q4 packed KV-cache quantization\n- reference quantized KV decode attention\n- experimental Metal q8/q4 GQA decode attention\n- GQA/MQA/MHA tests\n- memory/compression benchmarks\n \nOut of scope:\n \n- production KV quantization policy\n- paged quantized KV-cache\n- automatic runtime routing\n- model-quality claims\n \nCommands:\n```\npython examples/quantized_kv_decode_demo.py\npython benchmarks/bench_quantized_kv_decode_attention.py --bits 8 --B 1 --MAX_S 4096 --length 4096 --Hq 32 --Hkv 8 --D 128 --group-size 32 --backend all --validate --compare-fp16\npython benchmarks/bench_quantized_kv_decode_attention.py --bits 4 --B 1 --MAX_S 4096 --length 4096 --Hq 32 --Hkv 8 --D 128 --group-size 32 --backend all --validate --compare-fp16\n```\n\n### Paged quantized KV-cache\n\nThe repo includes experimental q8/q4 paged KV-cache decode attention. K/V pages can be stored in quantized form with per-page/token/head/group scales, and paged GQA/MQA decode attention dequantizes values while following the block table.\n\nCurrent scope:\n\n* q8 paged KV quantization\n* q4 packed paged KV quantization\n* reference paged quantized GQA decode attention\n* experimental Metal q8/q4 paged decode kernels\n* memory/compression accounting\n* benchmarks\n\nOut of scope:\n\n* production paged cache allocator\n* automatic runtime routing\n* paged sparse quantized Metal backend unless implemented\n* model-quality claims\n\nCommands:\n```\npython examples/paged_quantized_kv_demo.py\npython benchmarks/bench_paged_quantized_kv_decode.py --bits 8 --B 1 --MAX_S 4096 --length 4096 --PAGE_SIZE 16 --Hq 32 --Hkv 8 --D 128 --group-size 32 --backend all --validate --compare-fp16\npython benchmarks/bench_paged_quantized_kv_decode.py --bits 4 --B 1 --MAX_S 4096 --length 4096 --PAGE_SIZE 16 --Hq 32 --Hkv 8 --D 128 --group-size 32 --backend all --validate --compare-fp16\n```\n\n### Long-context runtime integration\n\nThe repo includes an experimental long-context runtime scaffold that combines prefix KV-cache reuse, sparse attention planning, KV offload/prefetch, and optional quantized-KV metadata into one explicit path.\n\nCurrent scope:\n\n- token-prefix reuse via `InMemoryPrefixCache`\n- sparse needed-position planning via `SparseAttentionPattern`\n- offload residency checks via `KVResidencyMap`\n- prefetch before attention via `ensure_blocks_ready_for_attention`\n- structured runtime reports via `LongContextRuntimeReport`\n- synthetic long-context demo\n- benchmark scaffold\n\nOut of scope:\n\n- production serving runtime\n- async flash scheduler\n- automatic runtime policy\n- real model quality\n- distributed prefix cache\n\nCommands:\n```\npython examples/long_context_runtime_demo.py\npython benchmarks/bench_long_context_runtime.py --prompt-len 512 --shared-prefix-len 384 --max-new-tokens 8 --prefix-reuse --sparse --window-size 128 --sink-tokens 4 --offload --offload-ratio 0.5 --num-layers 2 --hidden-size 512 --intermediate-size 2048 --num-heads 8 --num-kv-heads 2 --head-dim 64 --bits 4 --backend-preset fused_experimental\n```\n\n### Benchmarking and autotuning\n\n- Unified benchmark runner\n- Local benchmark report generator\n- Chip-specific backend registry\n- Local autotune cache\n\n## Project Goal\n\nMLX makes Apple Silicon a serious local machine learning platform, but high-performance transformer inference still benefits from custom fused kernels, layout-aware memory movement, quantized matvec, and backend choices tuned to the local machine.\n\nThis project investigates how to build and validate those primitives directly in MLX using custom Metal kernels.\n\nThe core workflow is:\n\n1. implement a pure MLX reference path\n2. add a correctness-first Metal backend\n3. test the Metal backend against reference\n4. add experimental optimized backends\n5. benchmark locally on Apple Silicon\n6. use explicit backend flags or autotuning to select kernels\n\nThe project is intentionally incremental: attention, decode, KV-cache, quantized matvec, MLP, and model-level scaffolding are developed as separate testable pieces before attempting larger fused runtime paths.\n\n## Current status\n\nThis repo is an experimental kernel research project, not a production inference runtime.\n\nCurrent capabilities include:\n\n- MLX custom Metal kernels with pure MLX reference implementations\n- attention, decode attention, and paged decode attention\n- KV-cache and paged KV-cache update helpers\n- q4/q8 dequantization and quantized matvec kernels\n- parallel and tiled q4/q8 decode matvec backends\n- RMSNorm, RoPE, SwiGLU, residual, and QKV layout helpers\n- fused and quantized decode-block composition\n- quantized MLP-block composition\n- explicit-only fused q4/q8 MLP experiments\n- toy transformer-layer decode benchmark\n- Llama-like config, weight layout, and model-adapter scaffolding\n- GQA/MQA reference and composed decode support\n- sparse and sliding-window GQA/MQA attention experiments\n- opt-in backend autotuning for local Apple Silicon machines\n\nStable defaults remain conservative. Experimental backends such as threadgroup, tiled, specialized, and simdgroup kernels are explicit-only unless enabled through documented flags or autotuning.\n\n## Results\n\nNo verified benchmark numbers are published from this checkout yet.\n\nThe local evidence file [`docs/performance_report_local.md`](docs/performance_report_local.md) records the exact commands attempted during the 2026-06-20 audit and why that run stopped before correctness or benchmark data could be produced. This workspace was audited on Windows 11, and `pip install mlx --break-system-packages` failed with `No matching distribution found for mlx`, so `pytest tests -q`, `python examples/run_basic.py`, the full benchmark suite, and autotune did not complete.\n\nUntil an Apple Silicon machine completes the commands below and produces raw JSON/CSV artifacts, this README intentionally does not publish a headline results table.\n\n```bash\npython benchmarks/run_all_benchmarks.py --full --output benchmarks/results/local_results.json --csv benchmarks/results/local_results.csv\npython scripts/save_benchmark_report.py benchmarks/results/local_results.json --output docs/performance_report_local.md\npython benchmarks/autotune.py --op all --quick --dtype float16 --write-cache\n```\n\n## Install\n\n```bash\npip install mlx pytest\npip install -e .\npytest tests -q\n```\n\nRun a basic example:\n\n```bash\npython examples/run_basic.py\n```\n\nRun the quick benchmark suite:\n\n```bash\npython benchmarks/run_all_benchmarks.py --quick\n```\n\nRun autotuning:\n\n```bash\npython benchmarks/autotune.py --op all --quick --dtype float16 --write-cache\npython benchmarks/run_all_benchmarks.py --quick --use-autotune\n```\n\n## Quick start\n\n```python\nimport mlx.core as mx\nfrom ops.attention_ops import fast_attention\n\nQ = mx.random.normal((1, 128, 8, 64)).astype(mx.float16)\nK = mx.random.normal((1, 128, 8, 64)).astype(mx.float16)\nV = mx.random.normal((1, 128, 8, 64)).astype(mx.float16)\n\nO = fast_attention(Q, K, V, causal=True, backend=\"auto\")\n```\n\n```python\nimport mlx.core as mx\nfrom ops.decode_ops import decode_attention\n\nq = mx.random.normal((1, 1, 8, 64)).astype(mx.float16)\nK_cache = mx.random.normal((1, 128, 8, 64)).astype(mx.float16)\nV_cache = mx.random.normal((1, 128, 8, 64)).astype(mx.float16)\n\ny = decode_attention(q, K_cache, V_cache, lengths=128, backend=\"metal\")\n```\n\n```python\nimport mlx.core as mx\nfrom ops.quant_ops import pack_q4, q4_matvec_decode\n\nx = mx.random.normal((1, 64)).astype(mx.float16)\nq4_vals = (mx.random.uniform((32, 64)) * 16).astype(mx.uint8)\npacked_w = pack_q4(q4_vals)\nscales = mx.ones((32, 2), dtype=mx.float32)\n\ny = q4_matvec_decode(x, packed_w, scales, group_size=32, backend=\"metal_tiled\")\n```\n\n## Benchmark examples\n\n### Attention\n\n```bash\npython benchmarks/bench_attention.py --backend all --S 128 --H 8 --D 64 --dtype float16\npython benchmarks/bench_threadgroup_attention.py --mode prefill --B 1 --S 128 --H 8 --D 64 --dtype float16 --backend all\npython benchmarks/bench_simdgroup_attention.py --mode prefill --B 1 --S 128 --H 8 --D 64 --dtype float16 --backend all\n```\n\n### Decode and KV-cache\n\n```bash\npython benchmarks/bench_decode_attention.py --B 2 --MAX_S 32 --H 8 --D 64 --length 32 --dtype float16 --backend all\npython benchmarks/bench_paged_decode_attention.py --B 2 --MAX_S 128 --PAGE_SIZE 16 --H 8 --D 64 --length 128 --dtype float16 --backend all\npython benchmarks/bench_decode_block.py --B 2 --MAX_S 128 --T 32 --H 8 --D 64 --dtype float16 --backend all\npython benchmarks/bench_gqa_decode_attention.py --B 1 --MAX_S 128 --Hq 32 --Hkv 8 --D 128 --dtype float16 --cache contiguous --backend all\npython benchmarks/bench_gqa_decode_block.py --B 1 --MAX_S 128 --Hq 32 --Hkv 8 --D 128 --T 16 --dtype float16 --cache contiguous\npython benchmarks/bench_gqa_attention.py --B 1 --S 128 --Hq 32 --Hkv 8 --D 128 --dtype float16 --causal --backend all --validate\n```\n\n## Sparse and sliding-window attention\n\nThe repo includes experimental GQA/MQA-aware sparse attention paths for long-context inference, including sliding-window attention and sink-token patterns.\n\nCurrent scope:\n\n- reference sparse GQA attention\n- sliding-window GQA prefill Metal kernel\n- sliding-window GQA decode Metal kernel\n- sink-token support\n- GQA/MQA/MHA tests\n- benchmarks against dense and sparse reference paths\n\nOut of scope:\n\n- production sparse-attention policy\n- automatic sparse routing\n- block-sparse optimized Metal backend unless implemented\n- model-quality claims\n\nCommands:\n\n```bash\npython benchmarks/bench_sparse_attention.py --B 1 --S 512 --Hq 32 --Hkv 8 --D 128 --window-size 128 --sink-tokens 4 --dtype float16 --backend all --validate\npython benchmarks/bench_sparse_decode_attention.py --B 1 --MAX_S 4096 --length 4096 --Hq 32 --Hkv 8 --D 128 --window-size 512 --sink-tokens 4 --dtype float16 --backend all --validate\n```\n\n### Quantization\n\n```bash\npython benchmarks/bench_dequant.py --bits 4 --M 4096 --K 4096 --dtype float16 --backend all\npython benchmarks/bench_quant_matvec_parallel.py --bits 4 --B 1 --K 4096 --N 4096 --group-size 32 --dtype float16 --backend all\npython benchmarks/bench_quant_matvec_tiled.py --bits 4 --B 1 --K 4096 --N 4096 --dtype float16 --backend all\n```\n\n### Transformer blocks\n\n```bash\npython benchmarks/bench_quantized_decode_block.py --bits 4 --cache contiguous --B 1 --K 4096 --H 32 --D 128 --MAX_S 128 --T 16 --dtype float16 --backend-preset parallel\npython benchmarks/bench_quantized_mlp_block.py --bits 4 --B 1 --S 1 --hidden-size 4096 --intermediate-size 11008 --dtype float16 --backend-preset all\npython benchmarks/bench_fused_mlp_kernels.py --bits 4 --B 1 --S 1 --hidden-size 4096 --intermediate-size 11008 --dtype float16 --backend-preset all --validate\npython benchmarks/bench_toy_transformer_decode.py --cache contiguous --bits 4 --B 1 --K 512 --H 8 --D 64 --INTERMEDIATE 1024 --MAX_S 64 --T 8 --dtype float16 --backend-preset parallel\npython benchmarks/bench_llama_layer_decode.py --bits 4 --B 1 --T 16 --hidden-size 512 --intermediate-size 2048 --num-heads 8 --num-kv-heads 2 --head-dim 64 --MAX_S 128 --dtype float16 --backend-preset all --validate\n```\n\nThe fused MLP path is experimental and explicit-only. It does not replace the stable composition-first `quantized_mlp_block` path or default backend routing.\n\nGQA prefill avoids expanding KV heads to query heads in the optimized path.\n\nThe full Llama-like decode layer experiment is a synthetic single-layer benchmark, not a full production model runtime.\n\n### Full benchmark suite\n\n```bash\npython benchmarks/run_all_benchmarks.py --quick\npython benchmarks/run_all_benchmarks.py --quick --use-autotune\npython benchmarks/run_all_benchmarks.py --full --output benchmarks/results/local_results.json --csv benchmarks/results/local_results.csv\npython scripts/save_benchmark_report.py benchmarks/results/local_results.json --output docs/performance_report_local.md\n```\n\n## Benchmark Suite\n\nThe repo includes a unified benchmark runner, machine metadata capture, CSV/JSON output, and a report-generation path for local Apple Silicon measurements.\n\nUseful entry points:\n\n```bash\npython benchmarks/run_all_benchmarks.py --quick\npython benchmarks/run_all_benchmarks.py --quick --use-autotune\npython benchmarks/run_all_benchmarks.py --full --output benchmarks/results/local_results.json --csv benchmarks/results/local_results.csv\npython benchmarks/autotune.py --op all --quick --dtype float16 --write-cache\n```\n\n## Autotuning\n\nAutotuning is opt-in. The repo includes a backend registry and a local JSON cache for machine-specific backend choices.\n\nExample:\n\n```python\nfrom ops.autotune_ops import select_backend\n\nbackend = select_backend(\n    \"decode_attention\",\n    {\"B\": 1, \"MAX_S\": 128, \"H\": 8, \"D\": 64, \"length\": 128},\n    \"float16\",\n    default_backend=\"metal\",\n)\n```\n\nCommands:\n\n```bash\npython benchmarks/autotune.py --op all --quick --dtype float16 --write-cache\npython benchmarks/run_all_benchmarks.py --quick --use-autotune\n```\n\nAutotune results are local to a specific machine and should not be treated as universal performance claims.\n\n## Llama-like integration scaffold\n\nThe repo includes lightweight model-integration scaffolding for Llama-like transformer layouts.\n\nCurrent scope:\n\n- Llama-like config objects\n- weight-shape mapping helpers\n- fused QKV layout conventions\n- model-adapter scaffold\n- toy and random decode examples\n- dependency-light checkpoint layout inspection and mapping helpers\n\nOut of scope for now:\n\n- production checkpoint loading\n- tokenizer integration\n- sampling loop\n- full model serving runtime\n\nExample commands:\n\n```bash\npython examples/inspect_model_shapes.py\npython examples/llama_like_decode_demo.py\n```\n\n## Checkpoint layout scaffold\n\nThe repo includes a dependency-light checkpoint layout scaffold for inspecting and validating Llama-like tensor names and shapes.\n\nCurrent scope:\n\n- JSON checkpoint manifests\n- Llama-style tensor name mapping\n- layer-shape validation against `LlamaLikeConfig`\n- fused QKV shape derivation\n- q4/q8 packaging shape specs\n\nOut of scope:\n\n- production safetensors loading\n- Hugging Face download support\n- tokenizer integration\n- full model execution from checkpoint\n\nCommands:\n\n```bash\npython examples/checkpoint_manifest_demo.py\npython examples/fuse_qkv_demo.py\n```\n\nThis scaffold is intended to prepare the repo for future checkpoint adapter work without adding heavy dependencies.\n\n## Real checkpoint adapter scaffold\n\nThe repo includes a dependency-light checkpoint adapter scaffold that connects local tensor stores or JSON manifests to the Llama-like config and kernel-facing weight layouts.\n\nCurrent scope:\n\n- in-memory tensor stores for tests and demos\n- manifest-backed shape-only tensor stores\n- optional safetensors store if `safetensors` is installed\n- Llama-style layer tensor validation\n- GQA-aware QKV shape handling\n- fused QKV creation from q/k/v tensors\n- quantized packaging specs for q4/q8 layouts\n\nOut of scope:\n\n- model downloads\n- tokenizer\n- sampling loop\n- production Llama or Mistral runtime\n- full checkpoint quantization or calibration\n\nCommands:\n\n```bash\npython examples/checkpoint_adapter_demo.py\npython examples/layer_weight_adapter_demo.py\n```\n\n```bash\npython examples/llama_layer_decode_demo.py\n```\n\n## Checkpoint-to-quantized packaging\n\nThe repo includes a correctness-first scaffold for packaging local floating-point layer tensors into the q4/q8 layouts expected by the custom matvec kernels.\n\nCurrent scope:\n\n- groupwise q4/q8 quantization for local tensors\n- q4 packing compatible with repo kernels\n- scale tensor generation\n- symmetric zero-point materialization for kernel compatibility\n- fused QKV quantization\n- layer-level `QuantizedLlamaLayerPackage`\n- conversion to `LlamaLayerKernelWeights`\n\nOut of scope:\n\n- GPTQ, AWQ, or SmoothQuant\n- production calibration\n- model-quality preserving quantization\n- full checkpoint conversion CLI\n- tokenizer and sampling runtime\n\nCommands:\n\n```bash\npython examples/quantize_layer_demo.py\npython examples/checkpoint_to_quantized_demo.py\n```\n\n## Multi-layer decode stack\n\nThe repo includes a correctness-first multi-layer decode stack scaffold for synthetic Llama-like experiments.\n\nCurrent scope:\n\n- explicit multi-layer composition over existing single-layer decode blocks\n- one KV-cache per layer\n- final RMSNorm plus optional `lm_head`\n- synthetic stack decode loop\n- synthetic multi-layer generation demo\n- GQA and MHA test coverage\n\nOut of scope:\n\n- production checkpoint execution\n- meaningful text generation from trained weights\n- optimized prefill stack\n- full production multi-layer runtime\n\nCommands:\n\n```bash\npython examples/llama_stack_decode_demo.py\npython examples/toy_multilayer_generation_demo.py\npython benchmarks/bench_llama_stack_decode.py --bits 4 --B 1 --T 16 --num-layers 2 --hidden-size 512 --intermediate-size 2048 --num-heads 8 --num-kv-heads 2 --head-dim 64 --MAX_S 128 --dtype float16 --backend-preset all --validate\n```\n\nThis is synthetic random-weight infrastructure for runtime plumbing and cache validation, not production text generation.\n\n## Full tiny-model generation demo\n\nThe repo includes a synthetic end-to-end generation demo that connects tokenizer, embeddings, multi-layer Llama-like decode stack, `lm_head` logits, sampling, and decoding.\n\nCurrent scope:\n\n- synthetic random weights\n- `CharTokenizer` by default\n- optional tokenizer adapters if installed\n- q4/q8 synthetic layers\n- greedy, top-k, and top-p sampling\n- benchmarkable generation loop\n\nOut of scope:\n\n- meaningful text generation\n- trained model checkpoint execution\n- production chat runtime\n- model downloads\n- tokenizer or chat-template correctness claims\n\nCommands:\n\n```bash\npython examples/full_tiny_generation_demo.py --prompt \"Hello\" --max-new-tokens 8 --greedy\npython examples/full_tiny_generation_with_package_demo.py\npython benchmarks/bench_tiny_generation_pipeline.py --prompt-len 8 --max-new-tokens 16 --num-layers 2 --hidden-size 512 --intermediate-size 2048 --num-heads 8 --num-kv-heads 2 --head-dim 64 --bits 4 --backend-preset fused_experimental --greedy\n```\n\nThe package-based demo is intentionally conservative today. The current quantized package format carries metadata and tensor-layout information, but not the tensor payloads needed for real loading, so the helper raises a clear fallback message and switches to synthetic weights.\n\n## Tokenizer/checkpoint package alignment\n\nThe repo includes structured validation for checking whether tokenizer metadata, model config, quantized package metadata, embeddings, `lm_head`, and generation pipeline settings agree before running a generation or smoke path.\n\nCurrent scope:\n\n- tokenizer vocab and special-token checks\n- config/package shape checks\n- q4/q8 quantization metadata checks\n- embedding and `lm_head` shape checks\n- `TinyGenerationPipeline` alignment validation\n- package inspector alignment mode\n\nOut of scope:\n\n- production tokenizer correctness\n- chat-template validation\n- trained-model quality\n- model downloads\n- full real-model runtime\n\nCommands:\n\n```bash\npython examples/alignment_demo.py\npython scripts/inspect_quantized_package.py /tmp/mlx_quant_package.json --validate-alignment\n```\n\n## Local real-model smoke test scaffold\n\nThe repo includes a local-only smoke-test scaffold for checking tokenizer, package metadata, tensor-data availability, alignment validation, and optional synthetic fallback generation.\n\nCurrent scope:\n\n- dry-run validation of local package JSON\n- optional local tokenizer validation\n- metadata-only package executability check\n- synthetic fallback generation\n- structured smoke-test reports\n- CLI script\n\nOut of scope:\n\n- model downloads\n- production inference\n- real trained checkpoint execution without tensor data\n- automatic Hugging Face loading\n- chat templates\n\nCommands:\n\n```bash\npython scripts/smoke_test_local_model.py --package /tmp/mlx_quant_package.json --dry-run\npython scripts/smoke_test_local_model.py --synthetic-fallback --no-dry-run --prompt \"Hello\" --max-new-tokens 4\npython examples/local_smoke_test_demo.py\n```\n\n## Optimized prefill stack\n\nThe repo includes a multi-layer prefill path that processes prompt tokens in sequence form, fills one KV-cache per layer, and then continues generation with decode.\n\nCurrent scope:\n\n- synthetic/random weights\n- contiguous KV-cache\n- GQA/MQA prefill attention\n- prefill-then-decode demo\n- benchmark against token-by-token decode ingestion\n\nOut of scope:\n\n- production prompt runtime\n- real trained checkpoint execution\n- paged prefill unless explicitly implemented\n- chat templates\n- model downloads\n\nCommands:\n\n```bash\npython examples/prefill_then_decode_demo.py\npython benchmarks/bench_llama_prefill_stack.py --bits 4 --B 1 --S 64 --num-layers 2 --hidden-size 512 --intermediate-size 2048 --num-heads 8 --num-kv-heads 2 --head-dim 64 --MAX_S 128 --dtype float16 --backend-preset all --validate\n```\n\n## Prefix KV-cache reuse\n\nThe repo includes a correctness-first prefix KV-cache reuse scaffold that caches previously computed KV-cache entries and reuses them when a prompt shares a prefix with a previously seen prompt.\n\nCurrent scope:\n\n- fingerprint-based cache safety (prevents cross-model/config reuse)\n- `InMemoryPrefixCache` with LRU eviction\n- `ops/kv_cache_reuse_ops.py` — clone, slice, copy, compare cache helpers\n- prefix matching (longest-common-token prefix)\n- exact match fast path (clone cached cache, no recomputation)\n- partial match suffix fallback via token-by-token decode\n- `TinyGenerationPipeline` integration via `use_prefix_cache=True`\n\nOut of scope:\n\n- paged cache reuse (reserved for future work)\n- production serving-cache\n- batch \u003e 1\n- automatic/background cache updates\n\nCommands:\n\n```bash\npython examples/prefix_cache_reuse_demo.py\npython benchmarks/bench_prefix_cache_reuse.py --validate --prompt-tokens 8 --reused-tokens 6\npytest tests/test_prefix_cache.py -v\npytest tests/test_kv_cache_reuse_ops.py -v\npytest tests/test_prefill_with_prefix_reuse.py -v\n```\n\n## Checkpoint converter scaffold\n\nThe repo includes a dependency-light checkpoint converter scaffold for turning local layer tensors into repo-native q4/q8 package metadata.\n\nCurrent scope:\n\n- JSON package metadata (QuantizedCheckpointPackage)\n- synthetic/in-memory conversion path\n- manifest dry-run planning\n- fused QKV packaging\n- q4/q8 layer package metadata\n- package inspection CLI\n- per-layer tensor metadata (original shapes, packed shapes, scales)\n- optional tensor-data writer (`.npy` payloads with SHA256 checksums)\n- synthetic tensor-data generation for testing\n\nOut of scope:\n\n- model downloads\n- tokenizer\n- production safetensors conversion\n- calibrated quantization\n- GPTQ/AWQ/SmoothQuant\n- full model serving\n\nCommands:\n\n```bash\npython scripts/convert_checkpoint.py --synthetic-demo --bits 4 --group-size 32 --output /tmp/mlx_quant_package.json\npython scripts/inspect_quantized_package.py /tmp/mlx_quant_package.json\npython examples/checkpoint_converter_demo.py\n```\n\n## Real tokenizer adapter scaffold\n\nThe repo includes optional tokenizer adapters for local tokenizer files.\n\nCurrent scope:\n\n- toy CharTokenizer and WhitespaceTokenizer\n- optional `tokenizers` JSON adapter (HFTokenizerAdapter)\n- optional SentencePiece `.model` adapter (SentencePieceTokenizerAdapter)\n- TokenizerAdapterFactory for auto-detection by file extension\n- describe_tokenizer helper for metadata\n- load_tokenizer_for_generation for generation scaffold integration\n- no network access\n- no model/tokenizer downloads\n- unified encode/decode interface for generation scaffolds\n\nOut of scope:\n\n- production chat formatting\n- Hugging Face hub downloads\n- transformers tokenizer auto-loading\n- prompt templates\n- model-specific tokenizer correctness guarantees\n\nCommands:\n\n```bash\npython examples/tokenizer_adapter_demo.py\npython examples/tokenizer_adapter_demo.py --tokenizer /path/to/tokenizer.json --kind hf-tokenizers\npython examples/tokenizer_adapter_demo.py --tokenizer /path/to/tokenizer.model --kind sentencepiece\n```\n\nNote: Optional tokenizer packages are not required for the repo. Install them separately only if you want local real-tokenizer experiments.\n\n## Tokenizer and sampling scaffold\n\nThe repo includes a lightweight tokenizer, sampling, and generation scaffold for synthetic single-layer experiments.\n\nCurrent scope:\n\n- toy character tokenizer\n- toy whitespace tokenizer\n- greedy, top-k, and top-p sampling utilities\n- repetition-penalty helper\n- synthetic single-layer generation loop\n- integration with the Llama-like layer decode scaffold when available\n- synthetic random weights only\n\nOut of scope:\n\n- production tokenizer support\n- BPE or SentencePiece\n- Hugging Face tokenizer loading\n- real text-quality generation\n- full multi-layer runtime\n- model downloads\n\nCommands:\n\n```bash\npython examples/sampling_demo.py\npython examples/toy_generation_demo.py\n```\n\nThe toy generation demo uses random synthetic weights and is intended to test plumbing, not language quality.\n\n## API examples\n\n### Attention API\n\n```python\nimport mlx.core as mx\nfrom ops.attention_ops import fast_attention\n\nQ = mx.random.normal((1, 128, 8, 64)).astype(mx.float16)\nK = mx.random.normal((1, 128, 8, 64)).astype(mx.float16)\nV = mx.random.normal((1, 128, 8, 64)).astype(mx.float16)\n\nO = fast_attention(Q, K, V, causal=True, backend=\"auto\")\nO_exp = fast_attention(Q, K, V, causal=True, backend=\"row_parallel\")\n```\n\n### Decode API\n\n```python\nimport mlx.core as mx\nfrom ops.decode_block_ops import decode_block_from_qkv\nfrom ops.gqa_ops import gqa_decode_block_from_qkv\n\ncos = mx.random.normal((132, 32)).astype(mx.float32)\nsin = mx.random.normal((132, 32)).astype(mx.float32)\npacked_qkv = mx.random.normal((1, 1, 3 * 8 * 64)).astype(mx.float16)\nK_cache = mx.zeros((1, 128, 8, 64), dtype=mx.float16)\nV_cache = mx.zeros((1, 128, 8, 64), dtype=mx.float16)\n\nout, K_cache, V_cache = decode_block_from_qkv(packed_qkv, K_cache, V_cache, cos, sin, 0, H=8, D=64, backend=\"auto\")\n\ngqa_qkv = mx.random.normal((1, 1, 32 * 128 + 2 * 8 * 128)).astype(mx.float16)\ngqa_K = mx.zeros((1, 128, 8, 128), dtype=mx.float16)\ngqa_V = mx.zeros((1, 128, 8, 128), dtype=mx.float16)\ngqa_out, gqa_K, gqa_V = gqa_decode_block_from_qkv(\n    gqa_qkv, gqa_K, gqa_V, cos, sin, 0, num_attention_heads=32, num_key_value_heads=8, head_dim=128, backend=\"reference\"\n)\n```\n\n### Quantized matvec API\n\n```python\nimport mlx.core as mx\nfrom ops.quant_ops import pack_q4, q4_matvec_decode\n\nx = mx.random.normal((1, 64)).astype(mx.float16)\nq4_vals = (mx.random.uniform((32, 64)) * 16).astype(mx.uint8)\npacked_w = pack_q4(q4_vals)\nscales = mx.ones((32, 2), dtype=mx.float32)\n\ny = q4_matvec_decode(x, packed_w, scales, group_size=32, backend=\"metal_tiled\")\n```\n\n### Toy transformer API\n\n```python\nimport mlx.core as mx\nfrom ops.toy_transformer_ops import make_toy_layer_weights, toy_transformer_decode_layer\n\nweights = make_toy_layer_weights(512, 1024, bits=4, group_size=32)\nx = mx.random.normal((1, 1, 512)).astype(mx.float16)\nK_cache = mx.zeros((1, 64, 8, 64), dtype=mx.float16)\nV_cache = mx.zeros((1, 64, 8, 64), dtype=mx.float16)\ncos = mx.random.normal((68, 32)).astype(mx.float32)\nsin = mx.random.normal((68, 32)).astype(mx.float32)\n\ny, K_cache, V_cache = toy_transformer_decode_layer(\n    x,\n    weights[\"attn_norm_weight\"].astype(mx.float16),\n    weights[\"ffn_norm_weight\"].astype(mx.float16),\n    weights[\"qkv_w\"],\n    weights[\"qkv_scales\"],\n    weights[\"out_w\"],\n    weights[\"out_scales\"],\n    weights[\"gate_w\"],\n    weights[\"gate_scales\"],\n    weights[\"up_w\"],\n    weights[\"up_scales\"],\n    weights[\"down_w\"],\n    weights[\"down_scales\"],\n    K_cache,\n    V_cache,\n    cos,\n    sin,\n    0,\n    bits=4,\n    H=8,\n    D=64,\n)\n```\n\n### Autotuning API\n\n```python\nfrom ops.autotune_ops import select_backend\n\nbackend = select_backend(\n    \"q4_matvec_decode\",\n    {\"B\": 1, \"K\": 4096, \"N\": 4096, \"group_size\": 32},\n    \"float16\",\n    default_backend=\"metal_parallel\",\n)\n```\n\n## Verification philosophy\n\nEvery optimized backend should have a pure MLX reference path.\n\nThe intended workflow is:\n\n1. generate random test inputs\n2. run reference implementation\n3. run Metal or custom backend\n4. compare outputs within dtype-appropriate tolerance\n5. only then benchmark\n\nExperimental backends may be disabled by default until they are repeatedly validated on Apple Silicon.\n\n## Roadmap\n\n- [x] Baseline MLX custom Metal attention\n- [x] Reference correctness path\n- [x] Row-parallel and tiled-K/V attention experiments\n- [x] RMSNorm, RoPE, SwiGLU, residual helpers\n- [x] KV-cache update and decode attention\n- [x] Paged KV-cache and paged decode attention\n- [x] Fused decode block helpers\n- [x] q4/q8 dequantization and decode matvec\n- [x] Parallel and tiled q4/q8 matvec\n- [x] Quantized decode block\n- [x] Threadgroup attention v2\n- [x] Simdgroup attention experiments\n- [x] Unified benchmark and report suite\n- [x] Chip-specific autotuning\n- [x] Toy transformer-layer decode benchmark\n- [x] Llama-like model integration scaffold\n- [x] Quantized MLP block\n- [x] GQA/MQA support\n- [x] Checkpoint layout loader scaffold\n- [x] Full tiny-model generation demo\n- [x] Optimized prefill stack\n- [x] Optimized GQA Metal decode attention\n- [x] Fused q4 MLP kernel\n- [x] Real checkpoint adapter scaffold\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmanishklach%2Fmlx-metal-kernels","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fmanishklach%2Fmlx-metal-kernels","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmanishklach%2Fmlx-metal-kernels/lists"}