Apple Silicon baseline (Asahi honeykrisp): no cooperative matrix caps GEMM at scalar path; capability map + tuning roadmap

[IlyaasK]

Environment

Apple M1 Pro (t6000/j316s, 10 cores, 32GB unified); Asahi Linux Fedora 44, kernel 7.1.5-400.asahi.fc44.aarch64; Mesa honeykrisp 26.1.5 (asahi_icd), Vulkan loader 1.4.341. Built with ./dev.sh, Fedora glslc/shaderc 2026.1.

What already works

• Build + generation smoke tests coherent (Qwen3-0.6B-bf16/8bit)
• bf16 runtime probe correctly falls back (no VK_KHR_shader_bfloat16 on driver) — MLX_VULKAN_DEBUG_CAPS shows extension_present=0 runtime_probe_supported=1
• Device classification fixed by goniz/mlx PR Serialize GPU benchmarks with host-shared flock #69 (device-name fallback; vendorID is Mesa's 0x10005, gen TPS improved 22.2→37.6 on short-context generate)

Capability ground truth (GPU0 Apple M1 Pro, from vulkaninfo)

Exposed: Vulkan 1.3 conformant; timelineSemaphore; bufferDeviceAddress; shaderInt8; subgroupSizeControl (but subgroup_min/max reported 0); shaderFloatControls2; VK_KHR_shader_integer_dot_product; VK_EXT_pipeline_robustness; VK_KHR_maintenance5; clustered subgroups; subgroup size 32; unified memory.

Absent: VK_KHR_cooperative_matrix; VK_NV_cooperative_matrix2; VK_KHR_shader_bfloat16; VK_EXT_shader_atomic_float (not in GPU0 list).

Note: llvmpipe (the CPU ICD also present) DOES expose coopmat — earlier confusion possible when grepping whole vulkaninfo output.

Root cause of prefill gap

cooperative_matrix_supported=0 inside MLX → prefer_matmul_coopmat1() false → every GEMM runs scalar matmul_direct_{f16,bf16}_aligned_f16acc (verified MLX_VULKAN_MATMUL_DEBUG=1). RDNA3 Strix Halo CI runs coopmat paths — hence prompt_tps 286 vs 4126.

Baseline numbers (./dev.sh benchmark, 5-trial avg, Qwen3-0.6B)

• bf16: prompt 285.934 tps, gen 24.395 tps, peak 2.092 GB
• 8bit: prompt 285.182 tps, gen 30.235 tps, peak 1.628 GB
• Strix Halo reference: bf16 4126/67.99, 8bit 4173/121.1

Tuning roadmap (impact order)

1. Scalar mul_mm AGX tile tuning (GpuArchitecture::Apple arm, MLX_VULKAN_MATMUL_SPEC sweeps)
2. fast.cpp flash-attn tuning arms for Apple
3. Integer dot product for quantized decode GEMV/GEMM (driver exposes it)
4. conv.cpp shmem padding config for Apple
5. Upstream: honeykrisp cooperative-matrix support would unlock the existing cm1 kernel family

[IlyaasK]

Decode dispatch characterization (follow-up measurements)

Quantized GEMV kernels are NOT the decode bottleneck on M1 Pro/honeykrisp:

• mx.quantized_matmul (2048x1024 int8 g64): ~46us/call streamed, ~2us hot single-shot — kernel itself is fast and already uses GL_EXT_integer_dot_product (unconditional in mul_mat_vecq.comp).
• Per-call graph+dispatch overhead outside compiled loops: ~200us/call flat regardless of matrix size — host-side, hidden by mx.compile.

MLX_VULKAN_DECODE_BATCH=1 end-to-end (Qwen3-0.6B, 4096-token prompt + 128 gen, 5-trial avg)

Model	Gen TPS off -> on	Peak mem off -> on
bf16	26.3 -> 32.7 (+24%)	2.06 -> 2.69 GB
8bit	32.8 -> 31.9 (~flat)	1.66 -> 2.87 GB

Mechanism: should_prefer_long_decode_recording() suppresses early submits during decode-like streams -> fewer vkQueueSubmit calls. honeykrisp's submit cost makes this valuable, but deferred buffering costs +0.6-1.2 GB peak memory and does not help the faster 8-bit decode cadence.

Conclusion: left opt-in via env; not safe to default-on globally (memory regression, 8bit neutral). Candidate follow-ups: adaptive enable keyed on model footprint vs working-set budget, or per-stream flush tuning (MLX_VULKAN_DECODE_MAX_RECORDED_OPS, MLX_VULKAN_MAX_DEFERRED_TOTAL_BYTES) to recover memory while keeping the submit reduction.

[IlyaasK]

Scalar mul_mm BK correctness analysis (for whoever fixes the tiny-BK path)

From reading mul_mm.comp / mul_mm_funcs.glsl while chasing the BK<=2 corruption:

• DATA_A_F32/DATA_A_F16 hardcode #define BK 32 and ignore specialization constant 3 entirely (mul_mm.comp:128-134); only quant/bf16/fp8 types consume the runtime BK.
• The bf16 scalar consumer loop (mul_mm.comp:320-360) is self-consistent: BK/BK_STEP(=2) iterations x VEC2 = BK elements per block, reading buf[row * SHMEM_STRIDE + i] with SHMEM_STRIDE = BK/2 + 1.
• Suspect area is the load side: loadr_a = x % (BK / LOAD_VEC_A / LOAD_VEC_BATCH_A) and loadstride_a = BLOCK_SIZE * LOAD_VEC_A * BATCH / BK (mul_mm.comp:210-216). For aligned bf16 (LOAD_VEC_A=2, BATCH=1) at BK=4: loadstride=64 > BM=32, so the l < BM copy loop runs a single iteration and column coverage depends entirely on invocation count — the layout math silently assumes BK in {8,16,32}. BK=2 additionally makes several of these expressions degenerate (BK/LOAD_VEC/BATCH -> 0 or 1).
• Measured boundary (bf16, M1 Pro): BK=4 correct & fastest legal point (1117 GFLOPS); BK=2 wrong results (rel err ~400) despite ~1677 GFLOPS; BM=64 and BLOCK_SIZE!=32 variants also corrupt — consistent with smem sizing/coverage assumptions rather than one localized typo.

Suggested fix shape: either static_assert/clamp spec'd BK to supported set with a loud error, or rework the copy-loop coverage math to be derived from BLOCK_SIZE and BK symmetrically. Happy to test patches on M1 Pro/honeykrisp.

[IlyaasK]

Coopmat readiness state (task #1 follow-up)

Verified how mlx-vulkan behaves across implementations available on this machine:

device	ext present	feature bit	our cooperative_matrix_supported	result
Apple M1 Pro (honeykrisp 26.1.5)	✗	—	false	scalar kernels (tuned in #70/#71)
llvmpipe (lavapipe, same install)	✓ rev2	0 (raw vkGetPhysicalDeviceFeatures2 probe)	false	scalar kernels

Conclusions:

• Our two-condition gate (extension present AND VkPhysicalDeviceCooperativeMatrixFeaturesKHR.cooperativeMatrix) is correct and conservative; llvmpipe demonstrates why checking the extension alone would be insufficient.
• No code change needed for coopmat readiness: when honeykrisp ships VK_KHR_cooperative_matrix with the feature bit set plus an f16 16x16x16 A/B config, prefer_matmul_coopmat1() and the _cm1 flash-attn/matmul shader families engage automatically.
• Verification command once a driver ships it: MLX_VULKAN_DEBUG_CAPS=1 should print ext_coopmat=1 feat_coopmat=1 coopmat=1; then MLX_VULKAN_MATMUL_DEBUG=1 should dispatch *_cm1 shaders.

[IlyaasK]

Decode deep-dive results (post #72 baseline)

Method: MLX_VULKAN_TRACE_BATCH=1 during full generate + controlled budget sweeps + benchmark confirmation.

Findings:

1. Submission cadence is no longer the bottleneck (post-#72). Traces show recordings flushing at rec_ops≈3 due to the 32 MB byte budget once KV writes accumulate (~330 MB cumulative recorded bytes/token), but raising the byte budget hurts (42.8 t/s @2 GB vs 48.6 default short-context) because longer recordings accumulate more hazard barriers.
2. Hazard barrier volume is the structural cost: ~59k overlapping-buffer-range barriers per 82-token run. Most are genuine RAW dependencies on the growing KV cache (attention reads the slice just written), plus scratch reuse (~830). These are inline pipeline barriers inside long recordings — correct, but they serialize shader scheduling on AGX.
3. Op-count-driven flushing is neutral-to-marginal: MLX_VULKAN_DECODE_MAX_RECORDED_OPS=512 gives 32.5 vs 32.0 tok/s benchmark gen (+1.3%, within cross-run variance). Not worth a default change.
4. Kernel side remains exonerated: quantized GEMV ~2 µs hot / 46 µs streamed; FA decode branch (n_rows==1) unaffected by prefill tuning.

Remaining decode levers, in impact order:

• Range-aware hazard elision for append-only KV growth: decode writes a disjoint KV slice per token; reads of prior slices need no barrier against older unsynced writes. Tracking per-range epoch instead of clearing whole read/write sets at each hazard could remove most of the 59k barriers. Requires careful WAW/WAR reasoning; suggest prototype behind an env flag.
• Compiled-loop op thinning: verify mlx-lm/omlx hot loops hit mx.compile; uncompiled graphs pay ~200 µs/primitive host cost that dwarfs GPU time.

No code change from this investigation; numbers above serve as the baseline any hazard-tracker rework should beat (bf16 gen 32.0 tok/s, peak 2.12 GB).

[IlyaasK]

Correction to "Decode deep-dive results": hazard rework is NOT the lever

Traced hazard types with MLX_VULKAN_TRACE_SYNC=1 during full generate:

61,883 hazard raw     <- read-after-write on consumed outputs
    0 hazard war
    0 hazard waw
2,438 action=barrier reason=overlapping-buffer-range

The proposed "range-aware hazard elision" was based on assuming disjoint KV-slice WAW/WAR noise; measurement shows zero such cases. Every barrier guards a genuine RAW edge (kernel N+1 consumes kernel N's output), which Vulkan requires as a memory dependency even inside one command buffer. These cannot be elided without changing what executes.

Revised remaining decode levers:

1. Fewer dependency edges: op fusion at the graph level (mx.compile coverage in mlx-lm/omlx hot loops) — reduces barrier count by construction.
2. Driver-side barrier cost: honeykrisp pipeline-barrier overhead per dispatch pair; only addressable upstream.
3. Everything else measured this session (submit cadence, byte budgets, op-count budgets) is second-order.

Apologies for the churn — baseline numbers in the previous comment stand; treat the "range-aware hazard elision" suggestion there as withdrawn.

[IlyaasK]

Served-throughput data point (omlx end-to-end)

Ran omlx serve (0.6.3rc2, OpenAI API) against this Vulkan build with Qwen3-0.6B-bf16:

• ~38-39 tok/s sustained generation (128-token completions, 3 runs: 3.25-3.36s wall each)
• Coherent output; first-request latency includes graph compilation warmup

Context: mlx-lm's generate_step runs fully eager in this fork (no mx.compile on _step; upstream ml-explore removed step-level compile as well), while omlx builds around compiled graphs. The served number landing above the eager-harness baseline confirms the earlier conclusion: host dispatch was a real tax, the serving stack already fuses what's fusable, and the residual decode ceiling is the genuine RAW barrier chain + kernel costs documented previously.

No further action identified at the device layer for decode.

[IlyaasK]

Prefill FA rewrite scope: PV path is the bottleneck (decomposition data)

Shape battery via mx.fast.scaled_dot_product_attention, Hq=16/Hkv=8, S=4096, bf16:

hsk	hsv	time	eff TFLOPS
128	128	210 ms	0.65
128	16	84 ms	0.92
16	128	455 ms	0.17
16	16	47 ms	0.36

Reading:

1. The QK^T/score side runs near 0.9 TFLOPS — no rewrite needed there.
2. The PV/V-accumulation side is the bottleneck: cutting it (A->B) removes >60% of runtime; making QK^T tiny while keeping wide V (C) is slower than baseline -> wide-V tiles hit a register/shared-memory pressure pathology in the scalar kernel.
3. Suspect region: flash_attn.comp:379-427 - V loads via kvsh[... d * D_split + d_tid] and Of[r][d] += Pf[r] * Vf. The D_split x thread mapping appears to produce poor access patterns / spills for HSV=128 on honeykrisp.

Measurement caveat before anyone chases numbers: identical configs measured 0.32 vs 0.65 TFLOPS across sessions/scripts today. Any rewrite claim requires an interleaved-median harness (round-robin configs within one process, median of N) - cold-clock single shots are not trustworthy on this machine.

Design sketch for an AGX-targeted variant (scalar path, since coopmat is unavailable):

• Keep the existing QK^T/online-softmax skeleton (proven fast enough).
• Restructure PV: wider contiguous V loads (u16vec8-class where legal), O accumulators pinned to registers across K-blocks, tile geometry chosen so each thread's V slice is stride-contiguous rather than D_split-interleaved.
• Numerics gate at every step vs composed f32 reference (harness exists).

Effort estimate: days. Expected ceiling if PV reaches QK^T-side efficiency (~0.9 T): prefill roughly 2x end-to-end on M1 Pro without waiting on upstream coopmat.

[IlyaasK]

Reading list for the FA PV-path rewrite (curated from wafer-ai/gpu-perf-engineering-resources)

In order, mapped to the decomposition findings above:

1. FlashAttention-2 - sec 3 work partitioning IS the fix for the wide-V pathology (per-warp K slices, smem O partials + single reduction, finer thread tiles for register occupancy). Design document for the variant.
2. Online normalizer calculation for softmax - audit current rescale ordering before rewriting; deferred-max division saves per-element vector ops per K-block.
3. CUTLASS GEMM measurement methodology - addresses the 0.32-vs-0.65 TFLOPS session variance; interleaved repeats + medians required for step-0 harness.
4. siboehm CUDA matmul - tiling progression validating the measured BK=4 / occupancy-first behavior on this driver.
5. Volkov SC08 - reason from measured hardware behavior; the session's working method in paper form.
6. HipKittens - prior art for subgroup-tiled attention without vendor tensor cores (same constraint class as our no-coopmat scalar path).

Everything CUDA-tensor-core-specific (FA3/FA4/TMA) and all serving-layer material does not transfer through glslc->NIR->AGX; validate any adopted pattern empirically on honeykrisp.

[goniz]

I've merged all the PRs and cleaned up a bit👍🏻