Added
- DeepSeek-V4-Flash / GLM sparse attention, ported with modifications from
omlx's `glm_moe_dsa` custom kernels (Apache-2.0; per-file OpenAI/Apple
notices preserved). `dsa_sparse_attention` runs the sliding local window plus
the indexer-selected pooled KV rows plus per-head attention sinks in one
flash-softmax dispatch (f32 accumulation); the lightning indexer is
`dsa_indexer_scores` (prefill GEMM), `dsa_indexer_score_decode` (fused
decode-width scores), and `dsa_topk_indices` (2-pass radix arg-select). All
accept `qL >= 1`, so decode, MTP verify (`qL = 2`), and prefill share them.
- `dsa_kv_qat` and `dsa_indexer_qat`: fused DeepSeek-V4 quantization-aware
round-trips - per-64-block FP8-E4M3FN on the main-attention KV, and a 128-wide
Hadamard transform then per-32-block FP4-E2M1 on the indexer activations -
each bit-identical to the equivalent MLX graph.
- `silu_limit` activation for the fused K-quant MoE GLU gather: the clamped
SwiGLU `silu(min(g, limit)) * clip(u, -limit, limit)` that DeepSeek-V4's
`LimitedSwiGLU` needs, passed as a constant-buffer limit (dead-arg-stripped
for the existing silu/gelu paths, so their kernels are unchanged).
- `moe_router_topk` gains `sqrtsoftplus` scoring (`sqrt(softplus(x))`),
score-plus-bias ranked selection, and an optional per-expert routed scale -
DeepSeek-V4 routing in one dispatch.
- `gather_qmm_seg` + `expert_tile_map`: expert-sorted MoE prefill as one GEMM
per expert segment instead of per-row gathers, with the 64-row tile map built
on the GPU from the sorted routing indices (no host sync). `KQuantSwitchLinear`
takes this arm on large prefill batches, gated by `KQ_SWITCH_GEMM_MIN_ROWS`
(default `512`). `nax_gather_enabled` reports whether the sorted-gather NAX
leaf is reachable, so the arm defers to it on tensor-unit GPUs.
- `docs/kernels.md`: a capability-grouped reference for the fused and
architecture-specific kernels beyond the four core codec ops.
- `sdpa_fa_verify` head_dim 512: 256-thread d-split kernel (gemma-4
global-attention verify/decode, folds to 32 rows) + vectorized K/V staging.
- `sdpa_fa_verify` now takes a 64-row query tile (up from 32), so a GQA-16
fold at `q_len` 4 stays on the matrix units instead of falling to the
stock materialized path. `q_len` 1 is also accepted now, routing plain
GQA decode through the same kernel.
- Non-NAX `gather_qmm_rhs`: steel simdgroup-mma GEMM for the sorted-rhs
(SwitchGLU prefill) gather leaf on GPUs without tensor units, all 19
codecs. Walks each row tile's per-expert segments and runs one full-K
matmul per segment, so the sorted batch no longer decomposes into
per-row `gather_qmv` calls. The row tile height adapts to the batch's
rows-per-expert (BM 16/32/64 at M/E thresholds 40/384): every segment pays a
full-tile mma pass, so a tile much taller than a segment wastes most of every
matmul. This is a large speedup on MoE prefill shapes at big batches, and more
at mid sizes where the adaptive tile kicks in.
`KQ_DISABLE_GATHER_RHS_ALU=1` forces the old per-row path;
`KQ_GATHER_RHS_BM` pins the tile height (retuning lever). On NAX machines
the NAX leaf still takes precedence; the new kernel serves the cases NAX
refuses (older macOS, `K % 64 != 0`, `KQ_DISABLE_NAX=1`).
Changed
- iq2_xxs / iq3_xxs Ext MoE gathers stage their dequant LUTs in threadgroup
memory, shared across the gather's lanes.
- q2_k / q3_k decode `qmm_t` weight loaders re-read each K-tile statelessly
instead of carrying a deep decoded-register cache; q6_k keeps its shallow
cache (its two-stream nibble decode is heavy enough that re-reading is
slower).
- `KQuantMultiLinear` memoizes its gather index arrays across calls, so a decode
loop reuses one constant pair instead of rebuilding the arange/broadcast graph
every layer.
Fixed
- `sdpa_fa_verify` read a lazily-strided query as packed: it trusted the
query layout at graph-build time, so a non-contiguous q tile was streamed
at the wrong stride. Layout is now checked at eval time.
- The sorted-rhs gather leaf mis-walked expert strides on the FIRST
evaluation of a lazily-sliced weight stack (e.g. a fresh `w[:, :n]` view):
the op trusted `w.flags().row_contiguous` at graph-build time, but flags of
an unevaluated array are meaningless, so the compaction copy was skipped
and the stride-less kernel read experts at the packed stride. Affected the
NAX leaf too. The op now always inserts the `Contiguous` node when the rhs
leaf is reachable (zero-copy at eval when already packed) and eval_gpu
additionally gates on `w.flags().row_contiguous`.