rlx-metal 0.2.11

Metal backend for RLX — Apple Silicon GPU via Metal Performance Shaders + custom MSL kernels
Documentation

rlx-metal

Apple GPU backend for RLX — MSL kernels + MPSGraph + ICB-batched dispatch. Two coexisting strategies:

  1. Thunk path (thunk/) — per-op MSL kernel + dispatch. Fine control, mature; default for now.
  2. MPSGraph path (mps_graph.rs + mps_graph_lower.rs) — lower subgraphs to MPSGraph and let Metal optimize the schedule. Opt-in per-op (e.g., RLX_MPSGRAPH_ATTENTION=1). Phase J extended this to Concat / FusedSwiGLU / RoPE cos-sin slice.

ICB (Indirect Command Buffer) batching (icb.rs) is the cross-cutting throughput unlock — Phase H made matmul-interleaved schedules use it.

What's here

  • MSL kernel library (kernels.rs, 2.2k LOC) — softmax, layer norm, RMS norm, attention, fused SwiGLU, RoPE, BERT-layer fusion. f16/f32 via a HalfFlag dispatch (Phase F). Phase I added f16 variants for rms_norm / softmax / reduce.
  • MPSGraph bridge — opt-in lowering to Apple's high-level graph compiler for the attention / concat / SwiGLU / RoPE-cos-sin paths.
  • MPS BLAS (mps_blas.rs) — descriptor-cached MPS matrix multiply.
  • ICB (Indirect Command Buffer) batching — segmented matmul schedules issue as one indirect dispatch instead of N command buffers.
  • thunk/ — Thunk enum + Op→Thunk lowering (ThunkSchedule split into submodules; see mod.rs).
  • backend/ — top-level Backend impl + execution; MetalExecutable decomposed into compile/run/encode/bind/set/read/output submodules (each a sibling impl MetalExecutable; see mod.rs).
  • calibrate.rs — measured GFLOP/s per kernel variant; cached in ~/.cache/rlx/metal-calib-<hwid>.json. Uses rlx_ir::Tick.
  • cost.rs — cost model that consumes calibration values.
  • device.rs / arena.rs — Metal device + buffer arena.
  • op_registry — MetalKernel trait + register_metal_kernel for downstream custom ops.
  • FFT — fft_gpu.msl pow-2 path: large N goes multi-kernel (bit-reverse gathers src→dst on-GPU — so src may be a preceding GPU op like rfft's Concat — then tiled inner + radix-4/2 outer stages); under native-gpu-fft, N ≤ 4096 uses an on-chip single-kernel radix-2/4/8 transform plus real-input (Concat([sig, zeros])) fusion. Op::Fft thunk / host fallback for f64 / C64 / non-pow2. MPSGraph skips graphs containing Op::Fft; fft_real subgraphs route through thunks automatically.
  • GGUF dequant (dequant_gguf.msl + backend::encode_dequant_gguf) — on-device dequant for every GGUF scheme (ids 0–23), including Q4_1 (21) and Q5_0 / Q5_1 (22 / 23). (id 21). IQ-family schemes consult a 33 KB grid LUT staged at session init (Kernels::iq_grid_buffer). Query runtime coverage with [backend::has_metal_dequant_kernel].
  • Fused decode GEMV — single-pass matvec when m == 1 (skips f32 scratch): Q4_K, Q4_0, Q4_1, Q8_0, IQ4NL, IQ2_XXS/XS/S, IQ3_XXS/S, IQ1_S/M (dequant_gguf.msl). See docs/gguf-backend-paths.md for shape constraints and disable env vars.
  • FP8 / NVFP4 block matmul — dequant_matmul_fp8 / dequant_matmul_nvfp4 MSL for non-GGUF QuantScheme::Fp8* / Nvfp4Block (CPU fallback when deferred host ops are pending).

Full backend matrix: docs/gguf-backend-paths.md. Parity vs rlx_gguf on real Qwen3-0.6B weights: tests/iq_full_real_weights.rs, tests/iq_mv_parity.rs, tests/iq4_dequant_parity.rs, tests/q8_q4_dequant_parity.rs.

Cargo features

Feature Description
native-splat (default) RLX-owned MSL tile raster (splat.msl) + CPU project/bin/sort via slang-splat-ref (Rust reference, no Slang compiler).

The crate is built unconditionally on macOS via rlx's metal feature (which enables native-splat); on other platforms it stubs out at link time.

Install

[dependencies]
rlx-metal = "0.1"

Or, more typically:

[dependencies]
rlx = { version = "0.1", features = ["metal"] }

Build / test

cargo build -p rlx-metal --release
cargo test  -p rlx-metal --release

Gating env vars worth knowing:

  • RLX_MPSGRAPH_ATTENTION=1 — opt into MPSGraph attention lowering (otherwise thunks).
  • RLX_VERBOSE=1 — calibration log.
  • RLX_METAL_DEQUANT_GPU_DISABLE=1 — CPU GGUF dequant instead of MSL.
  • RLX_METAL_Q4K_FUSED_DISABLE=1 / RLX_METAL_Q4K_SG_DISABLE=1 — disable fused Q4_K GEMV paths.
  • RLX_METAL_Q40_FUSED_DISABLE=1 / RLX_METAL_Q41_FUSED_DISABLE=1 / RLX_METAL_Q80_FUSED_DISABLE=1 — disable fused Q4_0 / Q4_1 / Q8_0 GEMV.
  • RLX_METAL_IQ4NL_FUSED_DISABLE=1 and RLX_METAL_IQ{2,3,1}*_FUSED_DISABLE=1 — disable fused IQ-family GEMV paths.

See docs/gguf-backend-paths.md for the full GGUF env table.

Status

Mature for the BERT / Nomic inference path used in burnembed. ICB matmul + MPSGraph attention are production. Tier-2 fused ops (FusedAttnBlock, FusedBertLayer) work; FusedNomicLayer is disabled pending a SwiGLU stride fix (see thunk.rs:3315).

Gotchas

  • Per-run cost is dominated by wait_until_completed (~150 µs); encoding cost is comparatively small. Fusing op chains into one command buffer is far more valuable than reducing kernel count.
  • Thunk::Attention only supports MaskKind::Custom (plan #20). The lowering asserts; non-Custom kinds are a future kernel addition. MPSGraph attention bails to thunks for non-Custom.
  • Don't trust microbenchmarks under thermal throttle. Run scripts/check-throttle.sh before measuring.
  • Phase G eliminated the f32↔f16 cast tax inside AutoMixedPrecision; follow-on work that adds new ops should respect the registry of natively-half kernels.

License

GPL-3.0-only.