aarch64: add runtime i8mm expert kernel#217
Draft
RonitBStudent wants to merge 2 commits into
Draft
Conversation
RonitBStudent
marked this pull request as ready for review
July 14, 2026 21:17
RonitBStudent
marked this pull request as draft
July 14, 2026 21:17
There was a problem hiding this comment.
Pull request overview
This PR adds a runtime-dispatched AArch64 i8mm/SMMLA “expert” matmul kernel for the existing int8-activation IDOT path, keeping the baseline AArch64 binary ISA unchanged while taking advantage of i8mm-capable CPUs.
Changes:
- Add i8mm feature detection + runtime dispatch (macOS
hw.optional.arm.FEAT_I8MM, LinuxAT_HWCAP2/HWCAP2_I8MM) and anI8MM=0kill-switch for same-binary A/B comparisons. - Implement a 2-output × 2-activation SMMLA tile (with packing + multiple accumulators) for both int8 and packed-int4 expert weights, including tail handling and an int4 S=1 short-I guard.
- Add dedicated correctness coverage for the i8mm tile path and a same-binary microbenchmark target (
i8mm-bench), plus documentation and clean integration.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| docs/ENVIRONMENT.md | Documents I8MM env var and its A/B kill-switch behavior. |
| c/tools/clean.py | Cleans new i8mm_bench build artifact. |
| c/tests/test_idot.c | Extends IDOT exactness tests to cover i8mm tile entrypoints and dispatched matmul parity. |
| c/tests/bench_i8mm.c | Adds same-binary NEON-vs-SMMLA microbenchmark for real expert shapes. |
| c/Makefile | Adds build/run targets for the new i8mm benchmark and updates .PHONY. |
| c/glm.c | Adds i8mm runtime detection, activation packing scratch, SMMLA 2×2 kernels, and dispatch integration into IDOT matmuls. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+25
to
+28
| uint8_t *w=malloc(wn); int8_t *x=malloc((size_t)S*I); | ||
| float *yn=malloc((size_t)S*O*sizeof(float)),*ym=malloc((size_t)S*O*sizeof(float)); | ||
| float *sc=malloc((size_t)O*sizeof(float)),*sx=malloc((size_t)S*sizeof(float)); | ||
| if(!w||!x||!yn||!ym||!sc||!sx){ fprintf(stderr,"bench_i8mm: OOM\n"); return 0; } |
Comment on lines
+36
to
+38
| if(memcmp(yn,ym,(size_t)S*O*sizeof(float))!=0){ | ||
| fprintf(stderr,"bench_i8mm: numeric mismatch: %s\n",label); return 0; | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add a runtime-dispatched AArch64 i8mm expert kernel for the existing int8-activation IDOT path.
SMMLAonly inside per-function i8mm targets, preserving the baseline AArch64 binaryhw.optional.arm.FEAT_I8MMon macOS andAT_HWCAP2/HWCAP2_I8MMon LinuxSMMLAlanes do useful workS, and odd output countsS=1, I<6144case on the existing NEON pathI8MM=0for same-binary A/B measurementsThis keeps the dependency-free default CPU path and file formats unchanged.
Why
The warm profile in #136 is matmul-bound on AArch64, while the current NEON IDOT path uses SDOT but not the available i8mm/SMMLA instructions. A direct instruction substitution was slower; the useful unit is a 2×2 tile with activation-pair packing and multiple independent accumulators.
Validation
Apple M2 Pro (12 cores, 32 GB), macOS arm64, Apple Clang 21.0.0. The benchmark was single-threaded because local libomp headers were unavailable. It uses the real GLM-5.2 expert shapes, one untimed call per path, 21 timed A/B pairs with alternating order, and reports medians.
Checks run:
make check OMPDIR=— clean CPU build, C suite, and 62 Python tests pass with zero compiler warnings./c/tests/test_idot— exact against scalar references for int8/int4 tiles, random full-range weights,-128edges, tails, andS=1..3/ odd-output dispatchglm.c,test_idot.c, andbench_i8mm.c— clean compile-out of the AArch64 pathgit diff --checkStill needed before ready
I do not have the full model snapshot locally. Please run the token-exact oracle (32/32 TF + 20/20 greedy) and
OMP_NUM_THREADS=1 make -C c i8mm-benchon an i8mm Linux machine/GB10 or Grace.I8MM=0provides the same-binary NEON control.Tracks the aarch64 i8mm frontier item in discussion #209.