feat(image): fused DiT kernels, a 3x3 VAE convolution and faster large-row NAX matmuls, on mlx 0.32.3 by asher · Pull Request #89 · asher/mlx-kquant

asher · 2026-10-11T18:06:12Z

No description provided.

…or tests

…he <5.13 ceiling

…limit

…s and a three-floor roofline

…probes at diffusion transformer shapes

…ness

…vs A aggregate and per-cell median CIs

… scalar loads and no-reuse bytes in side columns

… thermal curve for the cold rate

…ore calls per sample

…, f32 rows at M 1 and 2, and R1 and R2 calibration

…ithout NAX

…tches a known pattern

… single-threaded A/B runs, both off when unset

…ithout tile arms

…s of pooled pairs, and raw samples reach disk before the summaries

…regate rebuilds the reports from it, and the median CI has CPU tests up to 5000 samples

… round beyond one weight and one output rounding

…ve guard, and KQ_NAX_BM128_ODD_M=0 keeps the parity rule

…, and KQ_SPLITK_F32=0 keeps the activation-dtype partials

…ib loads

…Q_QMM_LOG logs

…ount, not by output bits

…rors other than a declined route

…d size per GPU core

…h block in float, and KQ_VERIFY_MMA_HALF=1 keeps the half kernels

…8 and 256, at K below 2048, on three codecs below 4096 rows and up to 16384 rows

…heir own test

…within 2 ulp

…batch row bound of the residual op

…mplate parameters of each codec kernel

…uantize then the stock matmul on m3 max with 40 gpu cores

…quantize then the stock matmul on m3 max with 40 gpu cores

…t nax

…e m5 inside its measured k and n range, with a vjp on the quantized weight and a lever to force or disable it

…s k and n range, vjp, lever and log line

…ed matmul raises as on the kernel, the route decision is one tested function, and its tests check large calls on row slices

…and a shorter changelog entry

…ernel's name, and its tests check every row in blocks

…uffer completes, and the changelog names the measured chip

…herent buffers, and its kernels need two simdgroups

… as the ci paravirtual gpu miscompiles coherent bfloat vector accesses

…y carries

… env builds, so it holds on gpus without nax

… they hold where mlx compile is off

…MLX_DISABLE_COMPILE holds for the tests after it

… nodes off too, as a gpu without nax builds them

…re that puts back the one found, so they hold under MLX_DISABLE_COMPILE

asher added 30 commits September 27, 2026 19:01

build: bump mlx pin to 0.32.2, nanobind 2.15.0, add pinned-mlx behavi…

6652a67

…or tests

build: name mlx 0.32.2 and nanobind 2.15.0 in install docs and CI

9ae7cf5

build: tools extra and CI exclude only transformers 5.13.0, lifting t…

e5b5666

…he <5.13 ceiling

chore(release): 0.5.0

0e493c0

feat(qmm): KQ_TUNE_LIVE reads the NAX tile levers per dispatch

256fbf0

feat(qmm): KQ_NAX_TILE and KQ_QMM_TILE pin the quantized matmul tile

cead1df

feat(ops): max_threadgroup_memory reports the GPU threadgroup memory …

bf27f00

…limit

feat(bench): diffusion transformer shape harness with paired tile arm…

1c0cd17

…s and a three-floor roofline

feat(bench): thermal curve, first-dispatch cost and float64 numerics …

05164e8

…probes at diffusion transformer shapes

fix(bench): discard the first call of each shape visit in the DiT har…

ebc731a

…ness

fix(bench): pairs follow the Thue-Morse bit, 3 inner calls, per-op A …

c80be78

…vs A aggregate and per-cell median CIs

fix(bench): roofline bytes floor counts compulsory traffic, with Q6_K…

95f732d

… scalar loads and no-reuse bytes in side columns

fix(bench): rename the first-round column to round 0 and point to the…

36f82e7

… thermal curve for the cold rate

fix(bench): evaluate a sample's calls together and give fast shapes m…

24855c6

…ore calls per sample

feat(bench): numerics probe adds a same-sign reduction, R2 references…

062d392

…, f32 rows at M 1 and 2, and R1 and R2 calibration

fix(bench): dense and affine tile labels follow MLX's rules on GPUs w…

33c42b3

…ithout NAX

test(qmm): every qmm_t tile kernel name outside split-K and gather ma…

fd84c62

…tches a known pattern

docs(qmm): KQ_QMM_TILE runs only 64,64,2,2 today, KQ_TUNE_LIVE is for…

5326837

… single-threaded A/B runs, both off when unset

fix(bench): modulation linears run at one row per image, M 1 and 2, w…

6ec382e

…ithout tile arms

fix(bench): median CI sums the binomial tail in integers for thousand…

a759bd5

…s of pooled pairs, and raw samples reach disk before the summaries

feat(bench): shape samples reach a raw log as each slot ends, --reagg…

d3e6f4f

…regate rebuilds the reports from it, and the median CI has CPU tests up to 5000 samples

feat(bench): route numerics probe shows which quantized matmul routes…

d8fc85c

… round beyond one weight and one output rounding

feat(qmm): BM=128 tile runs at every M from 4096 rows behind a row-li…

04dc76d

…ve guard, and KQ_NAX_BM128_ODD_M=0 keeps the parity rule

fix(qmm): split-K routes store float partials so the fold rounds once…

6b67928

…, and KQ_SPLITK_F32=0 keeps the activation-dtype partials

test(nax): take the bm128 codec list from a static set when no metall…

c13f685

…ib loads

docs(nax): mark the odd-M bm128 threshold provisional and list what K…

badf3d2

…Q_QMM_LOG logs

test(nax): tell split routes apart by their logged kernel and slice c…

dca04cf

…ount, not by output bits

fix(bench): probe the ALU split and verify_qmv routes and raise on er…

7b32994

…rors other than a declined route

feat(qmm): KQ_SPLITK_CAP caps the small-M split-K slice target by gri…

3c8ed7d

…d size per GPU core

fix(verify): verify_mma keeps activations in their dtype and sums eac…

626263b

…h block in float, and KQ_VERIFY_MMA_HALF=1 keeps the half kernels

asher added 28 commits October 6, 2026 22:25

docs(lora): the rank bound of the shape rule, and timings at ranks 12…

3996b7c

…8 and 256, at K below 2048, on three codecs below 4096 rows and up to 16384 rows

chore(merge): main 0.4.17 into the image kernel branch

f09532a

fix(bench): an env arm value may hold commas, as a tile override does

73cff7e

feat(nax): quantized matmul with a gated residual epilogue

e796c67

feat(nax): 3x3 convolution as an implicit gemm

4b54521

test(nax): the name enumerating tests leave the residual kernels to t…

22aeb1f

…heir own test

test(nax): hold each element of the residual and convolution kernels …

478ba63

…within 2 ulp

docs(kernels): the output channels where winograd is faster, and the …

5f381b0

…batch row bound of the residual op

refactor(kernels): the plain qmm_t tile and its simdgroup grid are te…

4cc7299

…mplate parameters of each codec kernel

feat(matmul): q4_k, q5_k and q6_k matmuls past a row floor run as deq…

621d94d

…uantize then the stock matmul on m3 max with 40 gpu cores

feat(matmul): nvfp4, ptq1_0 and q3_k matmuls from 2048 rows run as de…

436f4f6

…quantize then the stock matmul on m3 max with 40 gpu cores

docs(kernels): the dequantize rows of the plain matmul on gpus withou…

8c0f9f6

…t nax

docs(kernels): the measured gain of the dequantize rows on m3 max

99ebd07

fix(matmul): the dequantize route runs on every gpu without nax befor…

7c0d9e6

…e m5 inside its measured k and n range, with a vjp on the quantized weight and a lever to force or disable it

docs(kernels): the dequantize route on gpus without nax before m5, it…

6b6e0ee

…s k and n range, vjp, lever and log line

fix(matmul): a gradient wrt the weight or scales of a dequantize-rout…

e12089e

…ed matmul raises as on the kernel, the route decision is one tested function, and its tests check large calls on row slices

docs(kernels): the dequantize route's weight gradient, its glu bound …

a2c6be5

…and a shorter changelog entry

fix(matmul): a dequantize-routed matmul refuses jvp and vmap by the k…

8acad67

…ernel's name, and its tests check every row in blocks

docs(kernels): routed calls hold their dense copy until the command b…

820b411

…uffer completes, and the changelog names the measured chip

fix(hc): the fused hyper-connection cycle hands off through device-co…

fe5dbe3

…herent buffers, and its kernels need two simdgroups

fix(hc): the fused cycle passes h between threadgroups as uint2 bits,…

088d345

… as the ci paravirtual gpu miscompiles coherent bfloat vector accesses

build: bump mlx pin to 0.32.3 and nanobind 3.0.1

dc00ddf

chore(merge): main 0.4.18 into the 0.5 branch, whose hc fix it alread…

53cee00

…y carries

test(matmul): the jvp and vmap refusal test expects the glu node each…

2e94d4b

… env builds, so it holds on gpus without nax

test(glu): the peak bounds follow the stock activation's own peak, so…

2b2999c

… they hold where mlx compile is off

test(glu): the compile switch test puts back the switch it found, so …

f099606

…MLX_DISABLE_COMPILE holds for the tests after it

test(matmul): the refusal test's off case turns the lora and residual…

832f5c9

… nodes off too, as a gpu without nax builds them

test(compile): tests of mx.compile turn the switch on through a fixtu…

7e6812a

…re that puts back the one found, so they hold under MLX_DISABLE_COMPILE

asher merged commit 8ea5ff9 into main Oct 11, 2026
7 checks passed

asher deleted the kq-dit-nonnax branch October 11, 2026 20:24

Navigation Menu

Repository navigation

feat(image): fused DiT kernels, a 3x3 VAE convolution and faster large-row NAX matmuls, on mlx 0.32.3 - #89

feat(image): fused DiT kernels, a 3x3 VAE convolution and faster large-row NAX matmuls, on mlx 0.32.3#89
asher merged 167 commits into
mainfrom
kq-dit-nonnax

asher commented Oct 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

asher commented Oct 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant