Feature Request: Metal prefill speedup (~68%) on M5 Max (NAX / tensor API kernels) #30287

Prerequisites

• I am running the latest code. Mention the version if possible as well.
• I carefully followed the README.md.
• I searched using keywords relevant to my issue to make sure that I am creating a new issue that is not already open (or closed).
• I reviewed the Discussions, and have a new and useful enhancement to share.

Feature Description

Docs

• Summary: https://github.com/LXD-8/llama.cpp/blob/metal-m5-prefill/docs/backend/metal-m5-prefill-summary.md
• Full write-up (methodology, data, discarded approaches): https://github.com/LXD-8/llama.cpp/blob/metal-m5-prefill/docs/backend/metal-m5-prefill.md

Feature Description

On a MacBook Pro M5 Max (128 GB) the Metal backend prefills at roughly 60% of oMLX. The main reason is that the M5 NAX matrix units are mostly unused by the current ggml-metal kernels. I made 6 changes, all of them in the Metal backend only, each one switchable off with an environment variable. Compared with upstream b11541 (the latest tag at time of writing):

Decode speed is unchanged.

Per-change contribution (llama-bench, 8k context, pp2048):

These numbers come from different measurement setups and do not add up to the total. See section 5 of the full doc.

Provenance of the ideas: the NAX flash attention, the NAX mat-mat with split-K, the GDN lane change and the int8 mat-mat follow ideas from MLX and oMLX (steel_attention_nax, steel_gemm_fused_nax, gated_delta_pipelined, oq_a8), where such kernels are always on and have no switch; the f16 activation conversion and the SwiGLU fusion are my own. The kernels here are new ggml-metal implementations.

Related work in this repo

I searched issues, PRs and Discussions for M5, NAX, tensor API, prefill, split-K, GDN, SwiGLU and int8. No existing item proposes this set of changes as a whole. Closest neighbours, and how this work relates to them:

• metal: enable Metal 4.0 tensor API on M5+/A19+ #27461 (merged) and Compile bug: Metal tensor API silently inert on M5+/A19+ — MTLCompileOptions.languageVersion never set #27473: enable and repair the Metal 4.0 tensor API on M5+. My kernels run on that path; this request is performance work on top of it, not a fix.
• metal : add tensor API flash attention kernel for F16 KV #29570 (merged), tensor API flash attention for F16 KV: the closest overlap with my NAX flash attention. Mine is a register-resident kernel for head sizes 128/256, measured +7.9% on top of current master. If a maintainer judges it redundant, I will follow metal : add tensor API flash attention kernel for F16 KV #29570's author instead of opening a PR.
• metal : clamp K extent in tensor API mat-mat kernel #27450 (merged): correctness clamp on the K extent in the tensor mat-mat kernel. Mine adds a split-K path for long K, which is a throughput change, not a correctness fix.
• metal : skinny-batch mul_mm tiles + deterministic split-K for small-batch decode #27441 (open): also split-K, but for small-batch decode (8 < ne11 <= 32). Mine targets long-K prefill (large ne01), so the two regimes are complementary.
• metal : few-row MMA mat-mul and batched copies for speculative decoding #29869 (merged) and metal : few-row MMA kernels for MUL_MAT_ID and more src0 types #30047 (open): few-row MMA mat-mat for speculative decoding. Mine is int8 mat-mat for large batches.
• metal : tuning Flash-Attention kernel across M-chips #27668 and metal : add fa-vec tunings for M3 Max, M5 and M5 Pro #27863 (merged): per-chip FA tuning tables. Mine changes kernel structure instead of tuning parameters; I can post fa-vec rows for M5 Max in that thread separately.
• metal : wide query tile for kernel_flash_attn_ext #28439 (open): wide query tile on the SIMD-group FA path. Same goal, different mechanism.
• GDN: metal : add GDN kernel #20361 (merged), metal : gated_delta_net cache fusion #25788 (closed), metal: fused_gdn_ch (chunked prompt processing) needed for Qwen3.5 on older Apple GPUs #20342 (closed). My change is lane-occupancy tuning of the existing kernel, no API change.
• SwiGLU into mat-mat fusion: [metal] Add fused RMS_NORM + MUL + SWIGLU for Qwen3Next #16143 (open) fuses norm+mul+SwiGLU instead. The same direction exists on CUDA in CUDA: fuse FFN gate/up matmuls and GLU in MMQ with interleaved gate/up rows #29948 and CUDA: fuse ffn_gate + glu into the mul_mat_q epilogue #27341.
• No int8 prefill for Metal exists yet; the closest precedent is Vulkan int8 coopmat in vulkan: int8 coopmat1 matmul implementation for AMD RDNA3 and RDNA4 #27952 (merged).

AI usage disclosure: AI tools were used for research and for parts of the implementation and testing; I checked every result. I wrote the original text in Chinese; the English text above is a translation, which I back-translated and compared against the Chinese original below for consistency.

文档

• 摘要版: https://github.com/LXD-8/llama.cpp/blob/metal-m5-prefill/docs/backend/metal-m5-prefill-summary.md
• 完整版（测试方法、数据、被放弃的方案）: https://github.com/LXD-8/llama.cpp/blob/metal-m5-prefill/docs/backend/metal-m5-prefill.md

Feature Description

在 MacBook Pro M5 Max（128 GB）上，Metal 后端的预填充速度只有 oMLX 的约 60%。主要原因是 M5 的 NAX 矩阵单元没有被现有 ggml-metal kernel 充分利用。我在 Metal 后端内做了 6 项改动，每项都可以用一个环境变量单独关闭。与上游 b11541（撰写时的最新 tag）相比：

解码速度不变。

各项改动的贡献（llama-bench，8k 上下文，pp2048）：

这些数字的测法不同，不能相加，详见完整版文档第 5 节。

思路来源：NAX flash attention、NAX 矩阵乘（含长 K 的 split-K）、GDN lane 调整、int8 矩阵乘这四项借鉴了 MLX 和 oMLX 的思路（steel_attention_nax、steel_gemm_fused_nax、gated_delta_pipelined、oq_a8）；在这两个引擎里，对应的 kernel 常开、没有开关。f16 激活转换和 SwiGLU 融合这两项是我自己的。这里的 kernel 都是为 ggml-metal 新写的实现。

与社区已有工作的关系

我用 M5、NAX、tensor API、prefill、split-K、GDN、SwiGLU、int8 等关键词搜过 issues、PRs 和 Discussions，没有找到整体重复的条目。最接近的几项以及本工作与它们的关系：

• metal: enable Metal 4.0 tensor API on M5+/A19+ #27461（已合并）与 Compile bug: Metal tensor API silently inert on M5+/A19+ — MTLCompileOptions.languageVersion never set #27473：在 M5+ 上启用并修复 Metal 4.0 tensor API。我的 kernel 跑在这条路径上；本请求是之上的性能工作，不是 bug 修复。
• metal : add tensor API flash attention kernel for F16 KV #29570（已合并），F16 KV 的 tensor API flash attention：与我的 NAX flash attention 最重叠。我的实现是面向 head size 128/256 的 register-resident kernel，在当前 master 基础上测得 +7.9%。如果维护者认为重复，我会转而配合 metal : add tensor API flash attention kernel for F16 KV #29570 作者，不再另开 PR。
• metal : clamp K extent in tensor API mat-mat kernel #27450（已合并）：tensor 矩阵乘 kernel 中 K 维的正确性 clamp。我的改动是为长 K 增加 split-K 路径，属于吞吐优化，不是正确性修复。
• metal : skinny-batch mul_mm tiles + deterministic split-K for small-batch decode #27441（open）：也有 split-K，但针对小 batch 解码（8 < ne11 <= 32）。我的针对长 K 预填充（大 ne01），两者互补。
• metal : few-row MMA mat-mul and batched copies for speculative decoding #29869（已合并）与 metal : few-row MMA kernels for MUL_MAT_ID and more src0 types #30047（open）：面向投机解码的 few-row MMA 矩阵乘。我的是大 batch 的 int8 矩阵乘。
• metal : tuning Flash-Attention kernel across M-chips #27668 与 metal : add fa-vec tunings for M3 Max, M5 and M5 Pro #27863（已合并）：按芯片型号收集的 FA 调参表。我改的是 kernel 结构而不是调参参数；我也可以单独去该帖提交 M5 Max 的 fa-vec 数据。
• metal : wide query tile for kernel_flash_attn_ext #28439（open）：SIMD-group FA 路径的宽 query tile。目标相同，机制不同。
• GDN：metal : add GDN kernel #20361（已合并）、metal : gated_delta_net cache fusion #25788（closed）、metal: fused_gdn_ch (chunked prompt processing) needed for Qwen3.5 on older Apple GPUs #20342（closed）。我的改动是对现有 kernel 的 lane 占用调优，不改 API。
• SwiGLU 融合进矩阵乘：[metal] Add fused RMS_NORM + MUL + SWIGLU for Qwen3Next #16143（open）融合的是 norm+mul+SwiGLU，模式不同。CUDA 侧有同方向的 CUDA: fuse FFN gate/up matmuls and GLU in MMQ with interleaved gate/up rows #29948 与 CUDA: fuse ffn_gate + glu into the mul_mat_q epilogue #27341。
• Metal 上还没有 int8 预填充的先例；最接近的跨后端先例是 Vulkan 的 int8 coopmat vulkan: int8 coopmat1 matmul implementation for AMD RDNA3 and RDNA4 #27952（已合并）。

AI 使用说明：调研、部分实现和测试用了 AI 工具，所有结果都经过我人工检查。正文由我用中文先写好，上面的英文版由 AI 翻译；翻译结果已回译成中文，并与我的中文原文逐条核对，事实一致。

Motivation

• With all changes enabled, prefill roughly matches oMLX 0.7.1: +3% at 8k, -6% at 32k. Comparison setup: MTP and KV quantization disabled on both sides, oMLX int8 prefill enabled.
• Correctness: test-backend-ops passes for MUL_MAT, FLASH_ATTN_EXT and GATED_DELTA_NET. Perplexity is 3.5172 with all changes on and with all of them off.
• Hardware: MacBook Pro, M5 Max (40-core GPU, 128 GB), macOS 27.2, model Qwen3.8-27B Q4_K_M.

• 全部开启后，预填充与 oMLX 0.7.1 基本持平：8k 领先 3%，32k 落后 6%。对比条件：两边都关闭 MTP 和 KV 量化，oMLX 开启其 int8 预填充。
• 数值正确性：test-backend-ops 的 MUL_MAT、FLASH_ATTN_EXT、GATED_DELTA_NET 全部通过；困惑度全开、全关都是 3.5172。
• 测试设备：MacBook Pro（M5 Max，40 核 GPU，128 GB），macOS 27.2，模型 Qwen3.8-27B Q4_K_M。

Possible Implementation

Code is on the metal-m5-prefill branch of LXD-8/llama.cpp, one commit per change, each behind an environment variable toggle:
https://github.com/LXD-8/llama.cpp/tree/metal-m5-prefill

Per CONTRIBUTING.md I plan to split this into one PR per change. Before opening any PR I would like to know:

1. Is there interest in landing this work at all?
2. If yes, is there a preferred order or grouping for the individual PRs?

代码在 LXD-8/llama.cpp 的 metal-m5-prefill 分支，每项改动一个提交，每项都有对应的环境变量开关：
https://github.com/LXD-8/llama.cpp/tree/metal-m5-prefill

按 CONTRIBUTING.md 的要求，我计划按项拆成多个 PR。在开任何 PR 之前想先确认：

1. 这个方向是否有兴趣合入？
2. 如果有，各子 PR 是否有偏好的顺序或分组？