perf(qwen3.5): cut ~550 dispatches/cycle from the decode control flow (#168) Metal op-trace of one DFlash2 verify cycle showed ~2900 primitive dispatches; this drops it to ~2350 by removing redundant work, all behind existing numerical boundaries: - split_sections(): unequal Split views replace per-partition slice_axis copy dispatches at every hot-path split (GDN qkvz|ba, qkv|z, conv q|k|v tails, attention q|gate, MLP gate|up). - fused gating kernel gains EMIT_EXP so the per-step recurrence consumes exp(g) directly, removing an Exp dispatch per GDN layer. - GGUF f32 conv/norm sidecar params normalize to bf16 at load (matching the safetensors/reference numerics), removing the f32 activation promotion inside conv_general and fast_rms_norm. - affine quantized sidecars (scales/biases/bias) pre-cast f16->f32 at load: the call promotes to f32 either way, so this is bit-identical and removes two AsType dispatches per projection. - DFlash2 commit skips the GDN tape replay on full accept: the verify forward already advanced conv/recurrent state through all written rows, so replay_into was rebuilding identical values. - SDPA verify blocks with q_len*gqa > 32 (threadgroup limit of the fused vector kernel) split into two tail-aligned causal calls, keeping the fused kernel instead of the ~15-op decomposed path. - sigmoid-mul / SwiGLU sites route through the compiled fused kernels (attention out-gate, GDN conv SiLU, gated RMSNorm, DFlash2 MLP, MoE switch GLU). - GDN q/k inv_scale factors fold into constant RMSNorm weight vectors: one fused norm dispatch instead of norm + scalar mul. - GDN conv adopts the existing width-generic fused window_conv Metal kernel (history prepend + depthwise conv + SiLU + next-history in one dispatch) with the generic path kept as fallback. - vendored mlx gains MLX_METAL_OP_TRACE (1=counts, 2=names, 3=dtypes), the diagnostic used for the cycle accounting. Kill-switches: MLX_DISABLE_QWEN35_WINDOW_CONV, MLX_DISABLE_SDPA_VERIFY_SPLIT. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- _{Stack created with GitHub Stacks CLI • Give Feedback 💬} --- > [!NOTE] > **High Risk** > Changes touch speculative decode numerics (GDN fused kernels, SDPA split, dtype sidecars), a new compiled verify path that caches weights in process-global tapes, and KV/cache snapshot semantics—bugs could cause silent wrong tokens or stale graphs after reload. > > **Overview** > This PR targets **~550 fewer primitive dispatches per DFlash2 verify cycle** on Qwen3.5 by replacing copy-heavy splits, fusing common elementwise patterns, and hoisting dtype work to load time, while adding an optional **whole-forward `mlx::core::compile` path** for verify. > > **Graph and verify path:** New `compiled_graph` bridges Rust builders to MLX compile (trace once, replay with `shapeless` prefix lengths). `forward_dflash2_compiled` traces embedding → layers → LM head when `record_tape` and full-row logits apply; layers expose `forward_verify` with cache/RoPE as graph inputs. Compiled tapes are erased on model drop and after weight updates; failures latch `dflash2_compiled_verify_disabled` and fall back to eager. > > **Hot-path dispatch cuts:** `split_sections` does one Split primitive instead of multiple `slice_axis` copies (GDN, Q-gate/KV/MLP splits, conv tails). Causal SDPA for short verify blocks (`q_len·gqa > 32`) splits into two tail-aligned fused calls, with a numerical parity test. Attention/DFlash2/MoE use compiled **sigmoid×mul** and **SwiGLU**; decoder layers use fused **add+RMSNorm** when eligible. > > **GDN / Metal:** Optional fused `gdn_prepare` (conv + SiLU + q/k/v + norms + gating) and `window_conv`; gating kernel **`emit_exp`** and optional precomputed gates skip extra Exp ops. Q/k scale folds into `rms_norm_scaled` weights. > > **Load-time dtype:** GGUF f32 norm/conv sidecars cast to detected **compute dtype** via `sidecar_to_compute_dtype`; affine quant **scales/biases** promote f16→f32 once when activations aren’t f16. GDN verify snapshots clone recurrent state by default (`MLX_SNAPSHOT_STATE_COPY` restores copies). > > **DFlash2 extras:** Fused **top-16** Metal path for candidate selection (with argpartition fallback). Commit path documents that **MTP replay stays required** even on full accept for AR-exact recurrent state. > > ^{Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit
377e78f30616b5b78b5fca2dc9f7ba18c1782327. Bugbot is set up for automated
code reviews on this repo. Configure
[here](https://www.cursor.com/dashboard/bugbot).} --------- Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>