|
| 1 | +# MLX SSD Chunk-Region Fusion — honest scope (auto vs hand-call) |
| 2 | + |
| 3 | +Status: **MEASURED on Apple Metal** (this Mac, MLX fork 0.32.0.dev, python3.13, |
| 4 | +metal=True). CUDA build+measure on gb10 is a follow-up (no CUDA on this Mac). |
| 5 | + |
| 6 | +## TL;DR — what is and is NOT "mx.compile auto-fuses" |
| 7 | + |
| 8 | +- **NO-GO on the literal ask**: `mx.compile` (MLX's native C++ fuser) does **NOT** |
| 9 | + auto-recognize the SSD chunk region and does **NOT** emit one kernel for it. |
| 10 | + MLX's `compile.cpp` `is_fusable()` is an elementwise typeid allowlist |
| 11 | + (Add/Mul/Exp/Tanh/Select/Broadcast); a Matmul/Reduce/Scan node is a hard fusion |
| 12 | + boundary, and `Compiled::eval_gpu` only emits per-element loops. We did **not** |
| 13 | + modify `compile.cpp` — there is no general tensor-core GEMM/scan codegen inside |
| 14 | + `Compiled::eval_gpu`, and writing one is a separate large subsystem. |
| 15 | + - Direct evidence: `mx.compile(reference_chain)` traces **2564 nodes** and the |
| 16 | + fuser produces only elementwise `CompiledMultiplyExp` / |
| 17 | + `CompiledBroadcastMultiply...` Compiled nodes — **zero `CustomKernel`**. MLX |
| 18 | + by itself never routes to the mono SSD kernel. |
| 19 | + |
| 20 | +- **The reachable win (GO)**: a fusable-**REGION recognizer + proof-gated route** |
| 21 | + at the Python level — `mlx.ssd_region_fuse.compile_chunk_region()`. It |
| 22 | + RECOGNIZES the F0/F1/F2 chunk-scan composite (presented as a `ChunkScanRegion` |
| 23 | + marker), runs the z3 algebraic-equivalence proof + TLA/TLC race escalation, and |
| 24 | + on PROVEN **auto-substitutes** the whole region with ONE custom-primitive op |
| 25 | + (`ssd_chunk_scan_fused`, one `mx.fast.{cuda,metal}_kernel`). The collapse to ONE |
| 26 | + kernel is real and measured. The honest caveat: the **call site invokes |
| 27 | + `compile_chunk_region`**, not bare `mx.compile`; and the mono BODY is the |
| 28 | + hand-written wave8/`ssd_fused` kernel — the recognizer ROUTES, it does not |
| 29 | + synthesize tensor-core GEMMs from the graph. |
| 30 | + |
| 31 | +So: "compiler-style recognizer + proof-gated route to a prebuilt mono kernel" is |
| 32 | +TRUE and demonstrated; "mx.compile now synthesizes the fused tensor-core kernel |
| 33 | +by itself" would be FALSE. |
| 34 | + |
| 35 | +## Evidence (Apple Metal, bounded shapes batch=1, seqlen=128, nheads=2, |
| 36 | +## headdim=16, dstate=8, chunk=64; LOCAL 80GB guard respected) |
| 37 | + |
| 38 | +### Kernel count (graph export, unevaluated) |
| 39 | +- MONO (`ssd_chunk_scan_fused`): graph = exactly **ONE `CustomKernel`** node |
| 40 | + (7 inputs A–G, 2 outputs). This is the mono-fusion proof. |
| 41 | +- REFERENCE (unfused per-timestep recurrence): **4483** primitive nodes |
| 42 | + (ExpandDims 1024, Slice 640, Squeeze 770, Multiply 768, Exp 128, Broadcast 768, |
| 43 | + Add 256, Sum 128, Concatenate 1). The "6-kernel" SSD chain in production lowers |
| 44 | + to this many ops in the eager graph; the mono kernel collapses all of it to one |
| 45 | + launch with chunk state resident in threadgroup memory. |
| 46 | + |
| 47 | +### Parity (fp32, over all elements) |
| 48 | +- `ssd_chunk_scan_fused` vs reference recurrence: **max|dY| = 9.5e-7**, |
| 49 | + **max|d final_state| = 1.8e-7**. |
| 50 | +- Via the recognizer route `compile_chunk_region` (proven → mono): same |
| 51 | + **max|dY| = 9.5e-7**. |
| 52 | + |
| 53 | +### Timing (MEASURED, Metal, 50-iter mean after warmup) |
| 54 | +- mono = **0.92 ms**, reference (unfused) = **5.0 ms** → **5.47x**. |
| 55 | +- LABEL: MEASURED on Apple Metal. This is mono-custom-kernel vs the eager unfused |
| 56 | + recurrence, NOT vs a separately-tuned production 6-kernel CUDA path. The |
| 57 | + "vs 6-kernel" CUDA number is to be MEASURED on gb10 (EXTRAPOLATION until then). |
| 58 | + |
| 59 | +### Proof gate (z3 + TLA/TLC) — NON-VACUOUS |
| 60 | +- Correct F1 `cb = C @ B^T` rewrite → z3 **unsat** → `proved_by=z3` → route to mono. |
| 61 | +- Injected transposed-A bug (`_inject_transpose_bug=True`) → z3 **sat** → |
| 62 | + `proved_by=refuted` → substitution **BLOCKED + RAISE**. The gate refutes a wrong |
| 63 | + fusion (non-vacuous). Counter-witness reported: `operand_maps: COUNTER-WITNESS |
| 64 | + [k = 0, ...]`. |
| 65 | +- z3-unknown → TLA+/TLC bounded single-writer race obligation (tlc-bounded), via |
| 66 | + `verify_escalation.verify_with_escalation`. Prover genuinely uses z3 |
| 67 | + (`Solver(`, `unsat`, `sat` present in `_gemm_rewrite_proof`). |
| 68 | + |
| 69 | +### Control paths (RULE #1, all verified) |
| 70 | +- proven → route to mono. |
| 71 | +- z3-refuted → RAISE (substitution blocked). |
| 72 | +- unproven + forced (`MLX_SSD_REGION_FUSE_FORCE`) → RAISE. |
| 73 | +- unproven + auto → unfused reference chain (correct math, more kernels — NOT a |
| 74 | + degraded/zeroed fallback). Default OFF behind `MLX_SSD_REGION_FUSE`. |
| 75 | + |
| 76 | +## Build note |
| 77 | +Both `python/mlx/ssd_fused.py` and `python/mlx/ssd_region_fuse.py` are **pure |
| 78 | +Python** modules in the MLX package; our two commits touch **no C++** |
| 79 | +(`compile.cpp` untouched). The editable fork install (mlx 0.32.0.dev, |
| 80 | +cpython-313 `core.so`) loads them directly — **no MLX C++ rebuild required**, and |
| 81 | +the existing working build was not clobbered. |
| 82 | + |
| 83 | +## Files |
| 84 | +- `/Volumes/external/sources/mlx/python/mlx/ssd_fused.py` — mono custom-primitive |
| 85 | + kernel (Metal MSL + CUDA bodies), ONE `mx.fast.{metal,cuda}_kernel` op. |
| 86 | +- `/Volumes/external/sources/mlx/python/mlx/ssd_region_fuse.py` — region |
| 87 | + recognizer + proof-gated route (`ssd_chunk_region`, `compile_chunk_region`, |
| 88 | + `prove_region`, `_escalate_tlc`). |
| 89 | +- Prover: `cppmega_mlx/nn/_tilelang/_gemm_rewrite_proof.py` + |
| 90 | + `verify_escalation.py` (z3 + TLA/TLC escalation). |
0 commit comments