SIGN IN SIGN UP

Compile stateless MLP blocks for decode (#621)

Profiling the decode step showed it is op-count-bound: the per-layer
MLP/MoE chains cost ~10 ms of the ~17.5 ms step on Qwen3.6-35B, spread
over hundreds of tiny kernels. The blocks are stateless, so this PR
routes decode-shaped calls (up to 4 tokens) through per-module
`mx.compile` traces. No custom kernel. Opt-in via
`VLLM_METAL_COMPILED_MLP=1`.

M1 Ultra, kill switch off vs on:

|  | off | on |
|---|---:|---:|
| Qwen3.6-35B serve decode (tok/s) | 54.2 | 56.0–57.5 |
| `vllm bench latency --batch-size 1 --input-len 32 --output-len 256` |
7.41 s | 7.28 s |

---------

Signed-off-by: RickyChen / 陳昭儒 <ricky.chen@infinirc.com>
C
Chao-Ju Chen committed
4edfe6adda257d80e26c3367b15d5af472611a25
Parent: 150dd32
Committed by GitHub <noreply@github.com> on 8/18/2026, 7:58:05 AM