diff --git a/docs/benchmark_results/rocm-decode-profile-gfx1151-2026-09-30.md b/docs/benchmark_results/rocm-decode-profile-gfx1151-2026-09-30.md index 64343c86e..436fb2355 100644 --- a/docs/benchmark_results/rocm-decode-profile-gfx1151-2026-09-30.md +++ b/docs/benchmark_results/rocm-decode-profile-gfx1151-2026-09-30.md @@ -95,7 +95,7 @@ Which of those roles each port would actually take over depends on whether mlxce - **#2063** (`fused_add_rms_norm`, `fused_rope_qk_append`): both paths ship off on every backend (`FUSED_ADD_RMSNORM_DEFAULT` and `FUSED_ROPE_APPEND_DEFAULT` are `false`, `layers.rs:808` and `:820`, after #905 measured no decode win on Metal), and only `llama3.rs` (with `gemma.rs` and `iquestloopcoder.rs`, not profiled) calls them. With the opt-in on, Llama 3.1 would reach only the post-attention join (`llama3.rs:1201`): its `rope_scaling` builds a frequency table, which routes around the RoPE kernel (`llama3.rs:669`). - **#2064** (`gumbel_max_sample`, `rejection_sample`): greedy argmax calls neither. A sampled run would reach the whole sampler tail (both kernels are on by default where ported) except the few logit-bias dispatches `--ignore-eos` adds; the table counts the whole tail, so its #2064 figures are slightly high. -- **#2065** (`moe_gateup`, `moe_down`): reached by `qwen3_moe` (`qwen3_moe.rs:223`, single-token decode). Not by granite, whose `block_sparse_moe` calls `SwitchGLU::forward` and never `forward_fused_kernel`, and not by Nemotron-H, whose `fused_moe_forward` default branch is `gather_qmm`; its kernel branch needs `MLXCEL_FUSED_MOE_RELU2` and `moe_fc1_relu2` (#2069). +- **#2065** (`moe_gateup`, `moe_down`): reached by `qwen3_moe` (`qwen3_moe.rs:223`, single-token decode). At this profile's source (`929c80ab`) not by granite, whose `block_sparse_moe` called `SwitchGLU::forward` and never `forward_fused_kernel`; since issue #2245 (2026-10-11) granite's `block_sparse_moe` routes single-token decode through `forward_fused_kernel`, so the granite rows in the table above predate the wiring. Not by Nemotron-H, whose `fused_moe_forward` default branch is `gather_qmm`; its kernel branch needs `MLXCEL_FUSED_MOE_RELU2` and `moe_fc1_relu2` (#2069). - **#2067** (`ssm_update_kernel`): reached by both hybrids, since `ssm_kernel_available()` gates every single-token SSD step (`granitemoehybrid.rs:440`, `nemotron_h.rs:537` and `:635`). - **#2068** (paged attention): not on this path at all. The bench decodes one sequence into a dense `KVCache`; no paged kernel or paged graph fallback appears in any trace. diff --git a/src/lib/mlxcel-surgery/src/config_yaml_compat_tests.rs b/src/lib/mlxcel-surgery/src/config_yaml_compat_tests.rs index c099f12ea..e4548036a 100644 --- a/src/lib/mlxcel-surgery/src/config_yaml_compat_tests.rs +++ b/src/lib/mlxcel-surgery/src/config_yaml_compat_tests.rs @@ -55,8 +55,7 @@ fn non_finite_factors_reach_operation_validation() { let yaml = format!("version: 1\noperations: [{{op: scale, pattern: '*', factor: {value}}}]"); let error = parse_config_str(&yaml, None) - .err() - .expect("non-finite rejected") + .expect_err("non-finite rejected") .to_string(); assert!( error.contains("scale") && error.contains("finite"), @@ -72,8 +71,7 @@ fn non_finite_coefficients_reach_add_validation() { "version: 1\noperations: [{{op: add, pattern: '*', source: donor.safetensors, alpha: {value}}}]" ); let error = parse_config_str(&yaml, None) - .err() - .expect("non-finite rejected") + .expect_err("non-finite rejected") .to_string(); assert!( error.contains("(add)") && error.contains("alpha must be finite"), @@ -89,8 +87,7 @@ fn non_finite_ratios_reach_interpolate_validation() { "version: 1\noperations: [{{op: interpolate, pattern: '*', source_a: a.safetensors, source_b: b.safetensors, method: lerp, ratio: {value}}}]" ); let error = parse_config_str(&yaml, None) - .err() - .expect("non-finite rejected") + .expect_err("non-finite rejected") .to_string(); assert!( error.contains("(interpolate)") && error.contains("ratio must be a finite number"), diff --git a/src/models/granitemoehybrid.rs b/src/models/granitemoehybrid.rs index dd4fac418..3e5b4926b 100644 --- a/src/models/granitemoehybrid.rs +++ b/src/models/granitemoehybrid.rs @@ -747,7 +747,7 @@ impl GraniteMoeHybridAttention { // Dense SwiGLU MLP (`GraniteMoeHybridMLP`). -struct GraniteMoeHybridMLP { +pub(crate) struct GraniteMoeHybridMLP { gate_proj: UnifiedLinear, up_proj: UnifiedLinear, down_proj: UnifiedLinear, @@ -764,15 +764,15 @@ impl GraniteMoeHybridMLP { // Sparse MoE block (`block_sparse_moe`): softmax-over-top-k router + SwitchGLU. -struct GraniteMoeHybridMoE { - router: UnifiedLinear, - switch_mlp: SwitchGLU, +pub(crate) struct GraniteMoeHybridMoE { + pub(crate) router: UnifiedLinear, + pub(crate) switch_mlp: SwitchGLU, top_k: i32, num_experts: i32, } impl GraniteMoeHybridMoE { - fn forward(&self, x: &MlxArray) -> UniquePtr { + pub(crate) fn forward(&self, x: &MlxArray) -> UniquePtr { let orig_shape = mlxcel_core::array_shape(x); let hidden_dim = orig_shape[orig_shape.len() - 1]; @@ -797,8 +797,34 @@ impl GraniteMoeHybridMoE { let topk_logits = mlxcel_core::astype(&topk_logits, mlxcel_core::dtype::FLOAT32); let scores = mlxcel_core::softmax(&topk_logits, -1); - let expert_out = self.switch_mlp.forward(&x_flat, &topk_indices); - let result = moe_weighted_sum(&expert_out, &scores, mlxcel_core::array_dtype(&x_flat)); + // Apply routed experts. Fused single-token decode kernel (#268) on by + // default; MLXCEL_FUSED_MOE=0 forces the proven SwitchGLU + + // moe_weighted_sum path (also the automatic fallback when the kernel + // does not support the config). Granite routes through softmax over + // ONLY the top-k logits, which equals softmax-over-all then + // renormalize over the top-k (norm_topk_prob), so `scores` already + // carries the full combine weights the kernel needs; the launcher + // casts scores to f32 itself. The shared expert joins outside this + // block (the decoder layer sums it with the MoE output) and is + // unchanged (issue #2245). + let result = { + let fused = if mlxcel_core::array_shape(&x_flat)[0] == 1 + && crate::models::switch_layers::fused_moe_enabled() + { + self.switch_mlp + .forward_fused_kernel(&x_flat, &topk_indices, &scores) + .map(|out| mlxcel_core::reshape(&out, &[1, hidden_dim])) + } else { + None + }; + match fused { + Some(out) => out, + None => { + let expert_out = self.switch_mlp.forward(&x_flat, &topk_indices); + moe_weighted_sum(&expert_out, &scores, mlxcel_core::array_dtype(&x_flat)) + } + } + }; if orig_shape.len() > 2 { mlxcel_core::reshape(&result, &orig_shape) @@ -810,7 +836,7 @@ impl GraniteMoeHybridMoE { // Shared dense expert (`shared_mlp`): fused gate/up SwiGLU, always summed with MoE. -struct GraniteMoeHybridSharedMLP { +pub(crate) struct GraniteMoeHybridSharedMLP { input_linear: UnifiedLinear, output_linear: UnifiedLinear, shared_intermediate: i32, @@ -833,7 +859,7 @@ impl GraniteMoeHybridSharedMLP { // Per-layer feed-forward: dense MLP, or MoE + shared MLP. -enum FeedForward { +pub(crate) enum FeedForward { Dense(GraniteMoeHybridMLP), /// Dense feed-forward kept in the fused `shared_mlp` form (a single /// `input_linear` producing gate+up). Used for quantized dense checkpoints, @@ -1334,7 +1360,7 @@ impl GraniteMoeHybridModel { /// Build the per-layer feed-forward: dense `mlp`, or `block_sparse_moe` + /// `shared_mlp` when `use_moe`. -fn build_feed_forward( +pub(crate) fn build_feed_forward( weights: &WeightMap, args: &ModelArgs, prefix: &str, diff --git a/src/models/granitemoehybrid_tests.rs b/src/models/granitemoehybrid_tests.rs index d1d26f1eb..430670f45 100644 --- a/src/models/granitemoehybrid_tests.rs +++ b/src/models/granitemoehybrid_tests.rs @@ -26,7 +26,22 @@ //! `mlx-community/granite-4.0-h-350m-4bit` checkpoint and is not exercised here //! (no Metal device is assumed in unit tests). -use super::granitemoehybrid::ModelArgs; +use super::granitemoehybrid::{FeedForward, GraniteMoeHybridMoE, ModelArgs, build_feed_forward}; +use crate::models::switch_layers::{ + HONEST_EXPERT_GROUP_SIZE, insert_honest_affine_swiglu_experts, moe_weighted_sum, +}; +use mlxcel_core::weights::WeightMap; + +fn values(len: usize, start: f32, step: f32) -> Vec { + (0..len).map(|idx| start + idx as f32 * step).collect() +} + +fn insert_tensor(weights: &mut WeightMap, name: &str, values: Vec, shape: &[i32]) { + weights.insert( + name.to_string(), + mlxcel_core::from_slice_f32(&values, shape), + ); +} /// A trimmed `granite-4.0-h-350m` config (the dense-MLP hybrid). Values mirror /// the real checkpoint, including the full 32-entry `layer_types` interleave. @@ -198,6 +213,213 @@ fn granite_hybrid_quantization_read_from_config() { assert_eq!(bf16.bits(), 4); } +/// Hidden width of the fused-kernel geometry below: one quantization group per +/// input row, so one token carries 64 elements. Reuses the honest affine 4-bit +/// expert triple the Qwen3MoE fused-geometry tests use (issue #1884). +const FUSED_HIDDEN: i32 = HONEST_EXPERT_GROUP_SIZE; +const FUSED_EXPERTS: i32 = 2; + +/// A granite `GraniteMoeHybridMoE` over two honest affine 4-bit bf16 experts of +/// intermediate width `dff`, loaded through the family's own `build_feed_forward`, +/// so the fused-kernel decision under test runs on the exact type and stored +/// triple a converted granite-4.0-h checkpoint produces. +fn fused_moe_block() -> GraniteMoeHybridMoE { + let mut weights = WeightMap::new(); + insert_honest_affine_swiglu_experts( + &mut weights, + "model.layers.0.block_sparse_moe.switch_mlp", + FUSED_EXPERTS, + FUSED_HIDDEN, + FUSED_HIDDEN, + ); + // A non-quantized router and shared expert, so the declared pair can only + // reach the routed experts. + insert_tensor( + &mut weights, + "model.layers.0.block_sparse_moe.router.layer.weight", + values((FUSED_EXPERTS * FUSED_HIDDEN) as usize, -0.2, 0.004), + &[FUSED_EXPERTS, FUSED_HIDDEN], + ); + insert_tensor( + &mut weights, + "model.layers.0.shared_mlp.input_linear.weight", + values((2 * FUSED_HIDDEN * FUSED_HIDDEN) as usize, 0.01, 0.002), + &[2 * FUSED_HIDDEN, FUSED_HIDDEN], + ); + insert_tensor( + &mut weights, + "model.layers.0.shared_mlp.output_linear.weight", + values((FUSED_HIDDEN * FUSED_HIDDEN) as usize, 0.02, 0.002), + &[FUSED_HIDDEN, FUSED_HIDDEN], + ); + let args: ModelArgs = serde_json::from_value(serde_json::json!({ + "model_type": "granitemoehybrid", + "vocab_size": 8, + "hidden_size": FUSED_HIDDEN, + "intermediate_size": FUSED_HIDDEN, + "num_hidden_layers": 1, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "rms_norm_eps": 1e-5, + "embedding_multiplier": 1.0, + "attention_multiplier": 1.0, + "logits_scaling": 1.0, + "residual_multiplier": 1.0, + "layer_types": ["attention"], + "num_local_experts": FUSED_EXPERTS, + "num_experts_per_tok": FUSED_EXPERTS, + "shared_intermediate_size": FUSED_HIDDEN, + "quantization": { "group_size": 64, "bits": 4 }, + })) + .expect("fused-geometry granite config must parse"); + let ff = build_feed_forward(&weights, &args, "model.layers.0") + .expect("honest affine 4-bit granite experts must load"); + match ff { + FeedForward::Moe { moe, .. } => moe, + _ => panic!("num_local_experts 2 must build the MoE feed-forward"), + } +} + +/// Before issue #2245 granite's `block_sparse_moe` never called +/// `forward_fused_kernel`, so the fused single-token MoE kernel (#268) never +/// ran for the family and every decode step paid the `gather_qmm` fallback +/// (25.4% of decode GPU time on gfx1151, per the #2061 profile). The kernel +/// must decline nothing on an honest affine 4-bit expert triple of dff 64 and +/// compute the `forward` + `moe_weighted_sum` function within the tolerance +/// that test pins. It launches the kernel, which on a backend without a port +/// refuses across the bridge, so it runs only when +/// `fused_moe_kernels_available()` says both ports exist. +#[test] +fn granite_hybrid_fused_kernel_matches_the_gather_path_at_one_token() { + if !mlxcel_core::fused_moe_kernels_available() { + eprintln!("skipping: this backend has no fused MoE kernel port (#1803)"); + return; + } + let moe = fused_moe_block(); + let x = mlxcel_core::astype( + &mlxcel_core::from_slice_f32( + &values(FUSED_HIDDEN as usize, -0.4, 0.0125), + &[1, FUSED_HIDDEN], + ), + mlxcel_core::dtype::BFLOAT16, + ); + let indices = mlxcel_core::from_slice_i32(&[0, 1], &[1, FUSED_EXPERTS]); + let scores = mlxcel_core::from_slice_f32(&[0.25, 0.75], &[1, FUSED_EXPERTS]); + + let fused = moe + .switch_mlp + .forward_fused_kernel(&x, &indices, &scores) + .expect("honest affine 4-bit experts of dff 64 must dispatch the fused kernel"); + let fused = mlxcel_core::reshape(&fused, &[1, FUSED_HIDDEN]); + let gathered = moe_weighted_sum( + &moe.switch_mlp.forward(&x, &indices), + &scores, + mlxcel_core::dtype::BFLOAT16, + ); + // A structural check, not a numeric parity bound (that lives in + // `fused_moe_parity_tests`): swapped planes or a dropped expert land near + // 1, the bf16 quantized jitter between the two paths two orders lower. + let fused = mlxcel_core::utils::array_to_vec_f32(&mlxcel_core::astype( + &fused, + mlxcel_core::dtype::FLOAT32, + )); + let gathered = mlxcel_core::utils::array_to_vec_f32(&mlxcel_core::astype( + &gathered, + mlxcel_core::dtype::FLOAT32, + )); + let diff_sq: f64 = fused + .iter() + .zip(&gathered) + .map(|(a, b)| f64::from(a - b).powi(2)) + .sum(); + let ref_sq: f64 = gathered.iter().map(|b| f64::from(*b).powi(2)).sum(); + let nrms = (diff_sq / ref_sq.max(1e-20)).sqrt(); + assert!( + fused.iter().all(|v| v.is_finite()) && nrms < 0.1, + "the fused kernel must compute the gather path's function: nrms {nrms:e}" + ); +} + +/// Reproduction for issue #2245: before it, `GraniteMoeHybridMoE::forward` +/// always took the `SwitchGLU::forward` + `moe_weighted_sum` gather path, so +/// the fused single-token decode kernel never ran for the family. The +/// production `forward` must return, at one token, exactly the fused-kernel +/// computation (bit for bit, since it IS that computation) and not the gather +/// computation, which the two candidate values below pin. The candidates come +/// from the block's own router outputs, so the comparison is against the exact +/// indices and scores `forward` would hand the kernel. Runs only when both +/// fused MoE ports exist and the `MLXCEL_FUSED_MOE` gate is on, since the +/// gate is part of the routing under test. +#[test] +fn granite_hybrid_forward_routes_one_token_through_the_fused_kernel() { + use crate::models::switch_layers::fused_moe_enabled; + + if !mlxcel_core::fused_moe_kernels_available() { + eprintln!("skipping: this backend has no fused MoE kernel port (#1803)"); + return; + } + if !fused_moe_enabled() { + eprintln!("skipping: MLXCEL_FUSED_MOE disables the routing this test pins"); + return; + } + + let moe = fused_moe_block(); + let x = mlxcel_core::astype( + &mlxcel_core::from_slice_f32( + &values(FUSED_HIDDEN as usize, -0.4, 0.0125), + &[1, FUSED_HIDDEN], + ), + mlxcel_core::dtype::BFLOAT16, + ); + + // The router outputs `forward` computes, replicated with the block's own + // router so the candidates use the exact indices and scores the + // production step would. + let logits = moe.router.forward(&x); + let kth = 0; // two experts, top_k 2: every expert is routed + let indices = mlxcel_core::argpartition(&logits, kth, -1); + let indices_shape = mlxcel_core::array_shape(&indices); + let topk_indices = + mlxcel_core::slice(&indices, &[0, kth], &[indices_shape[0], indices_shape[1]]); + let topk_logits = mlxcel_core::take_along_axis(&logits, &topk_indices, -1); + let topk_logits = mlxcel_core::astype(&topk_logits, mlxcel_core::dtype::FLOAT32); + let scores = mlxcel_core::softmax(&topk_logits, -1); + + let fused = moe + .switch_mlp + .forward_fused_kernel(&x, &topk_indices, &scores) + .expect("honest affine 4-bit experts of dff 64 must dispatch the fused kernel"); + let fused_candidate = mlxcel_core::reshape(&fused, &[1, FUSED_HIDDEN]); + let gather_candidate = moe_weighted_sum( + &moe.switch_mlp.forward(&x, &topk_indices), + &scores, + mlxcel_core::dtype::BFLOAT16, + ); + + let routed = mlxcel_core::utils::array_to_vec_f32(&mlxcel_core::astype( + &moe.forward(&x), + mlxcel_core::dtype::FLOAT32, + )); + let fused = mlxcel_core::utils::array_to_vec_f32(&mlxcel_core::astype( + &fused_candidate, + mlxcel_core::dtype::FLOAT32, + )); + let gathered = mlxcel_core::utils::array_to_vec_f32(&mlxcel_core::astype( + &gather_candidate, + mlxcel_core::dtype::FLOAT32, + )); + + assert!( + routed.iter().zip(&fused).all(|(a, b)| a == b), + "one-token forward must take the fused kernel path, not the gather path" + ); + assert!( + routed.iter().zip(&gathered).any(|(a, b)| a != b), + "the gather candidate must differ from the fused path for this geometry, \ + so the bit-exact match above is a routing check" + ); +} + /// Regression for the `slice_update`-is-functional footgun: `inject_at_positions` /// MUST write the visual features into `h` at the image-token positions. A /// dropped `slice_update` return leaves the image slots zero and silently makes diff --git a/src/models/switch_layers.rs b/src/models/switch_layers.rs index 963b52755..c83f48148 100644 --- a/src/models/switch_layers.rs +++ b/src/models/switch_layers.rs @@ -1215,6 +1215,10 @@ impl SwitchGLU { /// non-SwiGLU activation (the kernel fuses `silu(gate) * up` and nothing /// else, so Kimi K3's SiTU experts stay on the gather path). Gated by the /// caller (`MLXCEL_FUSED_MOE`). + /// + /// Used by: AfMoE, BailingMoe, Cohere2Moe, Dbrx, Dots1, GraniteMoeHybrid, + /// Klear, Laguna, Lfm2, Mellum, MiniMax, Mixtral, OLMoE, PhiMoE, + /// Qwen2Moe, Qwen3Moe, and Qwen3VLMoE pub fn forward_fused_kernel( &self, x: &MlxArray,