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