feat(mlx): port FlashQLA to Apple Silicon#12

Open

0xClandestine wants to merge 11 commits into

QwenLM:mainfrom

0xClandestine:mlx-native-optimizations

Commits on Apr 30, 2026

feat: add flash_qla_mlx for Apple Silicon
0xClandestine
committed
test: add numerical correctness tests for flash_qla_mlx (Apple Silicon)
0xClandestine
committed

Commits on May 3, 2026

refactor: replace _kkt_solve in-place mutation with functional accumulation
0xClandestine
committed
perf: add mx.compile for fwd/bwd when cu_seqlens=None
0xClandestine
committed
perf: vectorize pack/unpack/prepare_chunk_offsets with gather ops
0xClandestine
committed
feat: expose chunk_size in public API, fix chunk_size bugs, add MLX benchmark
0xClandestine
committed
docs: document why parallel scan is not beneficial for inter-chunk recurrence
0xClandestine
committed
feat: add SDPA comparison to MLX benchmark
0xClandestine
committed
perf: reduce kkt_fwd einsum FLOPs for GQA by computing gram at Hk heads
0xClandestine
committed
perf: replace _kkt_solve Python loop with Metal kernel (2.2x speedup)
0xClandestine
committed
perf: eliminate KKT layout copies, cache triu masks, fix bwd list overhead
0xClandestine
committed