From 717afe664dcddc3fdb7d85552ca654b97a507203 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 2 Jul 2026 14:01:03 +0900 Subject: [PATCH] fix(models/deepseek-v3): correct MoE reshape for 3D forward input Follow-up hardening for #525 (round 2): the Kimi-VL port loads its real checkpoint but aborted during the forward pass with `[broadcast_shapes] Shapes (1,91,12288) and (1,91,6,1) cannot be broadcast`. The root cause is not in the MoonViT vision tower, connector, or merge; it is the shared DeepSeek-V3 MoE backbone that Kimi-VL reuses as its text model. `SwitchGLU::forward` ended with `reshape(&output, &[n_tokens, top_k, -1])`, where `n_tokens` and `top_k` were read from `indices.shape[0]` and `indices.shape[1]`. That is only valid for a 2D `[n_tokens, hidden]` token stream. The text stack runs on 3D `[batch, seq, hidden]` states (attention reads `batch`/`seq` from the leading axes), so `indices` is `[batch, seq, top_k]`; the reshape then read `n_tokens = batch = 1` and `top_k = seq = 91` and folded the real `top_k` (6 = num_experts_per_tok) and `hidden` (2048) axes into one, producing `[1, 91, 12288]`. `moe_weighted_sum` then broadcast that against `scores[..., None]` = `[1, 91, 6, 1]` and aborted. It surfaces via Kimi-VL because kimi-vl-a3b is the first small DeepSeek-V3-architecture MoE exercised end to end here (full DeepSeek-V3 is too large to run locally). Drop the rank-assuming reshape and return the already-squeezed `[..., top_k, hidden]`, matching the proven `deepseek_v2` sibling and the shared `moe_weighted_sum` contract. This is an identity for 2D input (the reshape was already a no-op there) and fixes the 3D path, so both the text-only and VLM embeddings paths are correct. Adds two deterministic MoE regression tests (no checkpoint, no GPU): `switch_glu_forward_preserves_rank_for_2d_and_3d` asserts the expert output keeps a distinct `top_k` axis for both 2D `[n, H]` and 3D `[B, L, H]` inputs, and `moe_block_3d_forward_matches_reference_shape` runs the full gate then SwitchGLU then moe_weighted_sum block on a 3D input and asserts a finite `[B, L, H]` result, reproducing the crash path. --- src/models/deepseek_v3.rs | 156 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 148 insertions(+), 8 deletions(-) diff --git a/src/models/deepseek_v3.rs b/src/models/deepseek_v3.rs index 36a47ec87..6556cfc8c 100644 --- a/src/models/deepseek_v3.rs +++ b/src/models/deepseek_v3.rs @@ -589,11 +589,11 @@ pub struct SwitchGLU { impl SwitchGLU { pub fn forward(&self, x: &MlxArray, indices: &MlxArray) -> UniquePtr { - let indices_shape = mlxcel_core::array_shape(indices); - let n_tokens = indices_shape[0]; - let top_k = indices_shape[1]; - - // Expand x for gather_qmm: [n_tokens, hidden] -> [n_tokens, 1, 1, hidden] + // `x` is `[..., hidden]` with one or more leading token axes: a flat + // `[n_tokens, hidden]` stream in the text-only path, or a 3D + // `[batch, seq, hidden]` tensor in the VLM embeddings path. The two + // `expand_dims` below insert the gather (`top_k`) and matmul axes for + // `gather_qmm` regardless of how many leading axes `x` carries. let x_exp = mlxcel_core::expand_dims(x, -2); let x_exp = mlxcel_core::expand_dims(&x_exp, -3); @@ -669,9 +669,15 @@ impl SwitchGLU { ) }; - // Squeeze: [n_tokens, top_k, 1, hidden] -> [n_tokens, top_k, hidden] - let output = mlxcel_core::squeeze_axis(&output, -2); - mlxcel_core::reshape(&output, &[n_tokens, top_k, -1]) + // Drop the singleton matmul axis, preserving every leading token axis: + // `[..., top_k, 1, hidden] -> [..., top_k, hidden]`. A previous + // `reshape(&[n_tokens, top_k, -1])` here assumed a 2D `[n_tokens, hidden]` + // input; for a 3D `[batch, seq, hidden]` input it misread `n_tokens` and + // `top_k` from the leading axes and folded the real `top_k` and `hidden` + // axes into one (`[batch, seq, top_k * hidden]`). That corrupted tensor + // then failed to broadcast against `scores[..., None]` in + // `moe_weighted_sum` (`[batch, seq, top_k*hidden]` vs `[batch, seq, top_k, 1]`). + mlxcel_core::squeeze_axis(&output, -2) } pub fn from_weights( @@ -1238,6 +1244,140 @@ mod tests { } } + // ---- MoE 3D-input regression (issue #525 round 2) --------------------- + // + // The Kimi-VL text backbone is this DeepSeek-V3 model, driven through + // `forward_impl_with_embeddings` on 3D `[batch, seq, hidden]` hidden + // states. `SwitchGLU::forward` used to end with + // `reshape(&output, &[n_tokens, top_k, -1])`, where `n_tokens`/`top_k` + // were read from `indices.shape[0]`/`[1]`. That is only correct for a 2D + // `[n_tokens, hidden]` input; for a 3D input it read the batch/seq axes as + // `n_tokens`/`top_k` and folded the real `top_k` and `hidden` axes into a + // single `top_k * hidden` axis. The following `moe_weighted_sum` then tried + // to broadcast `[batch, seq, top_k*hidden]` against `[batch, seq, top_k, 1]` + // and aborted (the reported `(1,91,12288)` vs `(1,91,6,1)` crash). + + /// Deterministic, varied, finite fill so quantization scales stay non-degenerate. + fn ramp(n: usize) -> Vec { + (0..n) + .map(|i| (((i * 37 + 11) % 97) as f32 - 48.0) * 0.01) + .collect() + } + + fn quantized_triple( + data: &[f32], + shape: &[i32], + group_size: i32, + bits: i32, + ) -> ( + UniquePtr, + UniquePtr, + UniquePtr, + ) { + let w = mlxcel_core::from_slice_f32(data, shape); + ( + mlxcel_core::quantize_weights_w(&w, group_size, bits), + mlxcel_core::quantize_weights_scales(&w, group_size, bits), + mlxcel_core::quantize_weights_biases(&w, group_size, bits), + ) + } + + /// Tiny quantized SwitchGLU: `gate`/`up` are `[E, inter, hidden]`, + /// `down` is `[E, hidden, inter]` (group_size 64, 4-bit). + fn tiny_switch_glu(e: i32, hidden: i32, inter: i32) -> SwitchGLU { + let (gs, bits) = (64i32, 4i32); + let (gate_weight, gate_scales, gate_biases) = quantized_triple( + &ramp((e * inter * hidden) as usize), + &[e, inter, hidden], + gs, + bits, + ); + let (up_weight, up_scales, up_biases) = quantized_triple( + &ramp((e * inter * hidden) as usize), + &[e, inter, hidden], + gs, + bits, + ); + let (down_weight, down_scales, down_biases) = quantized_triple( + &ramp((e * hidden * inter) as usize), + &[e, hidden, inter], + gs, + bits, + ); + SwitchGLU { + gate_weight, + gate_scales, + gate_biases, + up_weight, + up_scales, + up_biases, + down_weight, + down_scales, + down_biases, + group_size: gs, + bits, + } + } + + fn tiny_gate(e: i32, hidden: i32, top_k: i32) -> MoEGate { + MoEGate { + weight: mlxcel_core::from_slice_f32(&ramp((e * hidden) as usize), &[e, hidden]), + e_score_correction_bias: mlxcel_core::from_slice_f32(&vec![0.0f32; e as usize], &[e]), + top_k, + n_routed_experts: e, + routed_scaling_factor: 1.0, + n_group: 1, + topk_group: 1, + norm_topk_prob: true, + } + } + + #[test] + fn switch_glu_forward_preserves_rank_for_2d_and_3d() { + let (e, hidden, inter, top_k) = (4i32, 64i32, 64i32, 2i32); + let glu = tiny_switch_glu(e, hidden, inter); + let gate = tiny_gate(e, hidden, top_k); + + // 2D flat token stream: [n_tokens, hidden] -> [n_tokens, top_k, hidden]. + let x2 = mlxcel_core::from_slice_f32(&ramp((3 * hidden) as usize), &[3, hidden]); + let (idx2, _) = gate.forward(&x2); + let y2 = glu.forward(&x2, &idx2); + mlxcel_core::eval(&y2); + assert_eq!(mlxcel_core::array_shape(&y2), vec![3, top_k, hidden]); + + // 3D VLM embeddings path: [batch, seq, hidden] -> [batch, seq, top_k, hidden]. + // The pre-fix reshape collapsed this to [batch, seq, top_k*hidden]. + let x3 = mlxcel_core::from_slice_f32(&ramp((3 * hidden) as usize), &[1, 3, hidden]); + let (idx3, _) = gate.forward(&x3); + let y3 = glu.forward(&x3, &idx3); + mlxcel_core::eval(&y3); + assert_eq!(mlxcel_core::array_shape(&y3), vec![1, 3, top_k, hidden]); + } + + #[test] + fn moe_block_3d_forward_matches_reference_shape() { + // Full MoE block on a 3D `[batch, seq, hidden]` input exercises the exact + // crash path: gate -> SwitchGLU -> moe_weighted_sum. Before the fix this + // aborted on the broadcast mismatch; it must now return `[batch, seq, hidden]`. + let (e, hidden, inter, top_k) = (4i32, 64i32, 64i32, 2i32); + let moe = MoEBlock { + gate: tiny_gate(e, hidden, top_k), + switch_mlp: tiny_switch_glu(e, hidden, inter), + shared_experts: None, + }; + let x = mlxcel_core::from_slice_f32(&ramp((3 * hidden) as usize), &[1, 3, hidden]); + let out = moe.forward(&x); + mlxcel_core::eval(&out); + assert_eq!(mlxcel_core::array_shape(&out), vec![1, 3, hidden]); + + let mx = mlxcel_core::max_all(&out); + mlxcel_core::eval(&mx); + assert!( + mlxcel_core::item_f32(&mx).is_finite(), + "MoE output must be finite" + ); + } + /// Regression guard for the real Kimi-VL / Moonlight text config, whose /// `q_lora_rank` is JSON `null` (the backbone projects the query directly). /// Parsing this shape used to fail with "invalid type: null, expected