perf(models): generalize fused QKV+RMSNorm+RoPE to standard-RMSNorm QK-norm models #326

Summary

Generalize the existing fused QKV-projection + split + Q/K-norm + RoPE kernel so models using a standard RMSNorm (not only GemmaRMSNorm) can use it, then wire it into Qwen3 and Qwen3-MoE first, with a staged rollout to the other QK-norm families.

Source: idea.md "2. Generic Fused QKV + RMSNorm + RoPE For QK-Norm Models" (P0, expected payoff medium, risk medium-high).

Current state

• FusedQKVLinear::forward_split_norm_rope_quantized() exists but is specialized to GemmaRMSNorm and is documented as "Used by: Gemma3 dense attention path" (src/lib/mlxcel-core/src/layers.rs:1282).
• Qwen3 attention uses a fused QKV projection, then separate reshape, q_norm, k_norm, transpose, and RoPE ops (src/models/qwen3.rs, attention forward around line 115).
• Many models still apply Q/K RMSNorm as separate ops: Qwen3-MoE, Qwen3-VL, Qwen3-VL-MoE, OLMoE, OLMo2/3, Apertus, Exaone4, Exaone-MoE, and MiniMax.
• The decode-gap investigation reports dense Llama/Granite near parity, while the gap grows with extra small ops such as QK-norm and xIELU (docs/benchmark_results/moe-decode-gap-investigation.md).

Proposal

• Generalize the fused projection/split/norm/rope primitive so standard RMSNorm weights can use it, not only GemmaRMSNorm (for example, accept a norm trait/enum or a generic norm parameter).
• Wire it first into Qwen3 and Qwen3-MoE, which already have a fused QKV layout and good benchmark coverage.
• Then port compatible QK-norm families in small batches: Qwen3-VL text path, OLMoE/OLMo2/OLMo3, Apertus, Exaone4/Exaone-MoE, MiniMax. These follow-up family rollouts can be split into sub-issues once the primitive is proven.

Scope

• Generalize forward_split_norm_rope_quantized and update its // Used by: comment per the shared-function convention in docs/code-guidelines.md.
• Wire Qwen3 and Qwen3-MoE in this issue.

Validation

• Correctness: RMS < 5e-3 against the existing graph path, plus greedy parity for fixed prompts where deterministic parity is expected.
• Performance: mlxcel-bench-decode on qwen3-0.6b-4bit, qwen3-8b-4bit, qwen3-30b-a3b-4bit, and Apertus-8B-2509-4bit.
• Regression guard: include a non-QK-norm dense model (for example Llama) to prove common attention behavior did not change.
• Suggested first experiment (idea.md): prototype generic RMSNorm support on Qwen3 only first, to isolate the QK-norm hypothesis with minimal model spread.

Acceptance criteria

• The fused primitive accepts standard RMSNorm Q/K norms, not just GemmaRMSNorm.
• Qwen3 and Qwen3-MoE attention dispatch the generalized fused path on decode; the non-fused path remains as fallback.
• RMS < 5e-3 vs the graph reference; greedy temp-0 parity (or documented jitter) on the listed checkpoints.
• mlxcel-bench-decode numbers recorded; the non-QK-norm dense regression guard shows no change.
• // Used by: comment updated; INTEGRATION into the real Qwen3/Qwen3-MoE forward path verified.

Risk and priority

idea.md priority P0; risk medium-high. Directly targets the "extra small ops per layer" decode gap and yields one shared primitive reusable by many recently added QK-norm families.

References

• idea.md "2. Generic Fused QKV + RMSNorm + RoPE For QK-Norm Models".
• Code: src/lib/mlxcel-core/src/layers.rs (forward_split_norm_rope_quantized), src/models/qwen3.rs.
• docs/benchmark_results/moe-decode-gap-investigation.md, docs/code-guidelines.md (shared-function rule).