Repository navigation
fix(perf): use MLX fast RMSNorm in GatedDeltaNet decode path - #30
Merged
Merged
Conversation
GatedDeltaNet linear-attention layers (Qwen3.5, Qwen3-Next, KimiLinear) were normalizing Q/K and the gated output by expanding the RMSNorm formula into primitive ops on every decode step: `square` → `mean_axis` → `add(eps)` → `sqrt` → `divide` → `multiply(scale)`. Each of these is a separate MLX graph node, and at decode time the per-call overhead — not the FLOP count — dominates on Apple Silicon. Reference mlx-lm uses `mx.fast.rms_norm` for the same paths. This change replaces the expanded graphs with the fast kernel: - `src/models/gated_delta.rs` — add `scaled_fast_rms_norm_no_weight(x, scale, eps)` helper that wraps `mlxcel_core::fast_rms_norm_no_weight` and post-multiplies by a scalar (mirrors `mx.fast.rms_norm(x, None, eps) * scale`). Collapse `RMSNormGated::forward` to a single `fast_rms_norm` call before the SwiGLU gate. - `src/models/qwen3_5.rs`, `src/models/qwen3_next.rs`, `src/models/kimi_linear.rs` — replace the manual Q/K RMSNorm-and-scale blocks (three call sites per model: prefill, decode, and the speculative-decoding verify pass on Qwen3.5) with `scaled_fast_rms_norm_no_weight`. Q uses `inv_scale^2`, K uses `inv_scale`, where `inv_scale = head_k_dim^-0.5`, preserving the existing dtype-and-scale semantics. Measured on M5 Max with `qwen3.5-0.8b-4bit`, 100-token decode, `--profile --no-chat-template`, 5-run mean over the prompt `def fibonacci(n):\n `: - Before: 425.16 tok/s - After: 535.43 tok/s (+25.9%, 96% of mlx-lm's 555.43 tok/s on the same prompt) The other GatedDeltaNet families share the helper and benefit from the same kernel switch; their absolute numbers are tracked separately. Numerics are unchanged at the kernel level (same RMS formula, just compiled into one MLX primitive instead of six), and the speculative-decoding verify-pass path is updated alongside prefill/decode so draft/verify dtype agreement is preserved.
inureyes
added
status:review
inureyes
added a commit
that referenced
this pull request
May 18, 2026
This release focuses on Apple Silicon decode performance for VLM and linear-attention model families, plus a CLI fix for `mlxcel list`. Performance: - Gemma3n bf16 prefill path: gemma3n-e4b-bf16 decode 11.16 to 38.81 tok/s on M5 Max (#28). - Qwen3-VL text-only decode: QKV fusion + `forward_text_only` fast path. Qwen3-VL-30B-A3B-Instruct-4bit decode 56.00 to 146.29 tok/s (2.61x, 99% of mlx-lm parity); Qwen3-VL-32B-Instruct-4bit decode 18.79 to 27.33 tok/s (1.45x, 94% parity); image-in-prompt and DeepStack paths unchanged (#29). - GatedDeltaNet fast RMSNorm in Qwen3.5 / Qwen3-Next / KimiLinear: `mlx::fast::rms_norm` replaces the expanded square / mean / sqrt / divide graph for Q/K and gated output. qwen3.5-0.8b-4bit decode 425.16 to 535.43 tok/s on M5 Max (+25.9%, 96% parity); Qwen3.5 speculative-decoding verify pass updated alongside prefill / decode to preserve draft / verify dtype agreement (#30). Fixed: - Rewrite `mlxcel list` output to drive from the `ModelType` registry: fixes stale count, missing VLM family, missing ~30 model types, and removes a broken docs link (#27). Docs: - Refresh README performance snapshot and `docs/benchmarks.md` to match the latest M5 Max sweep.
inureyes
added
status:done
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
GatedDeltaNet linear-attention layers (Qwen3.5, Qwen3-Next, KimiLinear) were normalizing Q/K and the gated output by expanding the RMSNorm formula into primitive ops on every decode step: `square` → `mean_axis` → `add(eps)` → `sqrt` → `divide` → `multiply(scale)`. Each of these is a separate MLX graph node, and at decode time the per-call overhead — not the FLOP count — dominates on Apple Silicon. Reference mlx-lm uses `mx.fast.rms_norm` for the same paths. This change replaces the expanded graphs with the fast kernel: - `src/models/gated_delta.rs` — add `scaled_fast_rms_norm_no_weight(x, scale, eps)` helper that wraps `mlxcel_core::fast_rms_norm_no_weight` and post-multiplies by a scalar (mirrors `mx.fast.rms_norm(x, None, eps) * scale`). Collapse `RMSNormGated::forward` to a single `fast_rms_norm` call before the SwiGLU gate. - `src/models/qwen3_5.rs`, `src/models/qwen3_next.rs`, `src/models/kimi_linear.rs` — replace the manual Q/K RMSNorm-and-scale blocks (three call sites per model: prefill, decode, and the speculative-decoding verify pass on Qwen3.5) with `scaled_fast_rms_norm_no_weight`. Q uses `inv_scale^2`, K uses `inv_scale`, where `inv_scale = head_k_dim^-0.5`, preserving the existing dtype-and-scale semantics. Measured on M5 Max with `qwen3.5-0.8b-4bit`, 100-token decode, `--profile --no-chat-template`, 5-run mean over the prompt `def fibonacci(n):\n `: - Before: 425.16 tok/s - After: 535.43 tok/s (+25.9%, 96% of mlx-lm's 555.43 tok/s on the same prompt) The other GatedDeltaNet families share the helper and benefit from the same kernel switch; their absolute numbers are tracked separately. Numerics are unchanged at the kernel level (same RMS formula, just compiled into one MLX primitive instead of six), and the speculative-decoding verify-pass path is updated alongside prefill/decode so draft/verify dtype agreement is preserved.
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
This release focuses on Apple Silicon decode performance for VLM and linear-attention model families, plus a CLI fix for `mlxcel list`. Performance: - Gemma3n bf16 prefill path: gemma3n-e4b-bf16 decode 11.16 to 38.81 tok/s on M5 Max (lablup#28). - Qwen3-VL text-only decode: QKV fusion + `forward_text_only` fast path. Qwen3-VL-30B-A3B-Instruct-4bit decode 56.00 to 146.29 tok/s (2.61x, 99% of mlx-lm parity); Qwen3-VL-32B-Instruct-4bit decode 18.79 to 27.33 tok/s (1.45x, 94% parity); image-in-prompt and DeepStack paths unchanged (lablup#29). - GatedDeltaNet fast RMSNorm in Qwen3.5 / Qwen3-Next / KimiLinear: `mlx::fast::rms_norm` replaces the expanded square / mean / sqrt / divide graph for Q/K and gated output. qwen3.5-0.8b-4bit decode 425.16 to 535.43 tok/s on M5 Max (+25.9%, 96% parity); Qwen3.5 speculative-decoding verify pass updated alongside prefill / decode to preserve draft / verify dtype agreement (lablup#30). Fixed: - Rewrite `mlxcel list` output to drive from the `ModelType` registry: fixes stale count, missing VLM family, missing ~30 model types, and removes a broken docs link (lablup#27). Docs: - Refresh README performance snapshot and `docs/benchmarks.md` to match the latest M5 Max sweep.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
GatedDeltaNet linear-attention layers (Qwen3.5, Qwen3-Next, KimiLinear) were computing Q/K and gated-output RMSNorm by expanding the formula into primitive MLX ops on every decode step (
square→mean_axis→add(eps)→sqrt→divide→multiply(scale)). Six graph nodes per call; per-step overhead dominates on Apple Silicon. Reference mlx-lm usesmx.fast.rms_normfor these paths.This PR replaces those expanded graphs with the fast kernel.
Changes
src/models/gated_delta.rsscaled_fast_rms_norm_no_weight(x, scale, eps)— wrapsmlxcel_core::fast_rms_norm_no_weightand post-multiplies by a scalar (mirrorsmx.fast.rms_norm(x, None, eps) * scale).RMSNormGated::forwardto a singlefast_rms_normcall before the SwiGLU gate.src/models/qwen3_5.rs,src/models/qwen3_next.rs,src/models/kimi_linear.rsscaled_fast_rms_norm_no_weight.inv_scale^2, K usesinv_scale, whereinv_scale = head_k_dim^-0.5— same numerics as before, one MLX primitive instead of six.Performance (M5 Max, 4-bit, decode tok/s)
Measured 5-run mean on
def fibonacci(n):\n(100 tokens,--profile --no-chat-template). Before runs: 427.17, 428.86, 423.32, 423.65, 422.78. After runs: 535.44, 535.26, 535.06, 536.12, 535.25.Other GatedDeltaNet families (Qwen3-Next, KimiLinear) share the helper and benefit from the same kernel switch; their absolute numbers are tracked separately.
Test plan
cargo fmt --all -- --checkcargo check --features metal,accelerateqwen3.5-0.8b-4bit100-tok decode reproduces ~535 tok/s ondef fibonacci(n):\n.