Skip to content

metal: speed up IQ2_XXS MoE prefill with half LUT and split MPP - #864

Open
mobailabs wants to merge 1 commit into
antirez:mainfrom
mobailabs:main
Open

metal: speed up IQ2_XXS MoE prefill with half LUT and split MPP#864
mobailabs wants to merge 1 commit into
antirez:mainfrom
mobailabs:main

Conversation

@mobailabs

@mobailabs mobailabs commented Aug 25, 2026

Copy link
Copy Markdown

Summary

  • Rewrite IQ2_XXS mul_mm_id dequant to a baked half LUT (0.25 * grid in binary16), then apply sign by XOR on the half bits.
  • Rewrite kernel_mul_mm_id_mpp in place as split-tile (two M=16 matmul2ds) + double-buffer pipeline + tail cull. MPP threadgroup is 12 KiB. Pipeline names are unchanged.
  • Q2_K f16 down shares this MPP kernel, so it picks up the same split/pipeline path.

Test plan

  • Same-machine A/B on MacBook Pro M5 Max 128 GB, Metal, DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf
  • Official ds4-bench sweep (promessi_sposi.txt, ctx 2048–65536 step 2048, 128 greedy tokens)
  • Two optimized runs vs one origin run at c1d4597
  • make ds4_test && ./ds4_test --metal-kernels on the PR tree

Prefill is about +5–8% vs the same-machine origin baseline (32-point mean about +6.4%). Decode is unchanged.

Context Prefill origin → opt Prefill Gen origin → opt
2048 778.8 → 835.5 +7.3% 45.38 → 45.32
8192 669.6 → 725.1 +8.3% 41.49 → 41.31
16384 527.9 → 561.5 +6.4% 39.74 → 39.73
32768 443.8 → 477.2 +7.5% 31.97 → 32.59
65536 334.7 → 350.7 +4.8% 27.37 → 27.12

…IQ2_XXS dequant into a half LUT and rewrite mul_mm_id_mpp as split-tile + pipeline + tail cull (12 KiB threadgroup). On M5 Max Flash IQ2 this is about +5-8% prefill vs the same-machine baseline;decode is unchanged. Q2_K f16 down shares the same MPP kernel.
GiorgioOppo added a commit to GiorgioOppo/ds4 that referenced this pull request Aug 25, 2026
Use an exact binary16 dequantization LUT and split the Metal MPP routed matmul into double-buffered 16-row tiles with tail culling. Preserve the packed route map integration from this branch.

Adapted from antirez#864.

Co-authored-by: cole <imzhpe@qq.com>
@ivanfioravanti

Copy link
Copy Markdown
Contributor

Great boost here! Wow!

GiorgioOppo added a commit to GiorgioOppo/ds4 that referenced this pull request Sep 1, 2026
Use an exact binary16 dequantization LUT and split the Metal MPP routed matmul into double-buffered 16-row tiles with tail culling. Preserve the packed route map integration from this branch.

Adapted from antirez#864.

Co-authored-by: cole <imzhpe@qq.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants