Skip to content

fix(models/deepseek-v3): correct MoE reshape for 3D forward input - #614

Merged
inureyes merged 2 commits into
mainfrom
fix/issue-525-deepseek-v3-moe-3d-reshape
Jul 2, 2026
Merged

fix(models/deepseek-v3): correct MoE reshape for 3D forward input#614
inureyes merged 2 commits into
mainfrom
fix/issue-525-deepseek-v3-moe-3d-reshape

Conversation

@inureyes

@inureyes inureyes commented Jul 2, 2026

Copy link
Copy Markdown
Member

Follow-up hardening for #525 (round 2): the Kimi-VL port loads its real checkpoint (round-1 config fix #608) but aborted during the forward pass with [broadcast_shapes] Shapes (1,91,12288) and (1,91,6,1) cannot be broadcast (running -p "Describe this image in one sentence." --image tests/fixtures/test_image.png, after expanding 1 media placeholder to 64 image tokens).

Root cause

The crash is not in the MoonViT vision tower, connector, or merge; those produce a correct [1, 91, 2048] merged-embedding stream. It is the shared DeepSeek-V3 MoE backbone that Kimi-VL reuses as its text model.

SwitchGLU::forward in src/models/deepseek_v3.rs 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. Note 12288 = 6 * 2048 and 91 = 64 image tokens + prompt. 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).

What changed

  • src/models/deepseek_v3.rs SwitchGLU::forward: drop the rank-assuming final 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. deepseek_v32 has its own already-correct SwitchGLU and is unaffected.
  • Adds two deterministic MoE regression tests (no checkpoint, no GPU): switch_glu_forward_preserves_rank_for_2d_and_3d and moe_block_3d_forward_matches_reference_shape.

Test plan

  • cargo test --lib deepseek_v3 (17 pass, including the two new tests)
  • cargo test --lib kimi_vl (16 pass, no regression)
  • cargo clippy --lib --tests -- -D warnings (clean)
  • cargo fmt

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.
@inureyes inureyes added type:bug Bug fixes, error corrections, or issue resolutions priority:high High priority status:review Under review labels Jul 2, 2026
@inureyes
inureyes merged commit 78b4df2 into main Jul 2, 2026
5 checks passed
@inureyes
inureyes deleted the fix/issue-525-deepseek-v3-moe-3d-reshape branch July 2, 2026 05:22
@inureyes inureyes self-assigned this Aug 31, 2026
@inureyes inureyes added status:done Completed and removed status:review Under review labels Sep 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

priority:high High priority status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant