Skip to content

[Metal:Perf] Fused-Q4 M8 tile and K-split for speculative-block decode - #4772

Merged
wangzhaode merged 1 commit into
masterfrom
sync/ali-to-github
Aug 14, 2026
Merged

[Metal:Perf] Fused-Q4 M8 tile and K-split for speculative-block decode#4772
wangzhaode merged 1 commit into
masterfrom
sync/ali-to-github

Conversation

@wangzhaode

Copy link
Copy Markdown
Collaborator

Automated sync from the internal MNN development repository.

  • Pending commits: 1
  • Head branch: sync/ali-to-github
  • Commit authors are preserved from the source commits.

Commits

The fused Q4 GEMM stage kernel uses an M32 matmul2d tile. A speculative
block is 8 rows, so it runs at quarter occupancy, and for projections
with few output tiles it also dispatches too few threadgroups to
saturate bandwidth.

Add two variants behind FUSED_Q4_REAL_UNPACK:
  - conv1x1_fused_q4_gemm_stage_m8: descriptor sized to M=8.
  - conv1x1_fused_q4_gemm_stage_ksplit{,_m8} + _ksplit_reduce: split K
    four ways across gid.z into fp32 partials, then sum + bias +
    activation in a second pass.

Both are gated on KVMeta::spec_block so nothing outside speculative
decoding changes path.

New switches, all default to auto (unset):
  MNN_METAL_FUSED_Q4_KSPLIT   =1 force on, =0 force off
  MNN_METAL_FUSED_Q4_M8       =0 force off (=1 same as unset, the tile is
  MNN_METAL_FUSED_Q4_KSPLIT_M8    only correct for area <= 8)

K-split reassociates the K accumulation, so output is not
bit-identical: on an iPad tier-9 sweep 8 of 9 prompts match the
baseline verbatim and one diverges on a near-tie greedy argmax.

The area-gate part of this work already landed with the fused-Q4
spec-block gate.

Discussed-in: Merge-Request 29266048 , URL: https://code.alibaba-inc.com/AliNN/AliNNPrivate/codereview/29266048
GitOrigin-RevId: c465bb4195f3722161c2338440b301f43e052b01
@wangzhaode
wangzhaode merged commit fe05d70 into master Aug 14, 2026
12 of 13 checks passed
@wangzhaode
wangzhaode deleted the sync/ali-to-github branch August 14, 2026 08:15
@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

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.

3 participants