fix(models/deepseek-v3): correct MoE reshape for 3D forward input - #614
Merged
Conversation
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.
…k-v3-moe-3d-reshape
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.
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::forwardinsrc/models/deepseek_v3.rsended withreshape(&output, &[n_tokens, top_k, -1]), wheren_tokensandtop_kwere read fromindices.shape[0]andindices.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 readsbatch/seqfrom the leading axes), soindicesis[batch, seq, top_k]; the reshape then readn_tokens = batch = 1andtop_k = seq = 91and folded the realtop_k(6 = num_experts_per_tok) andhidden(2048) axes into one, producing[1, 91, 12288].moe_weighted_sumthen broadcast that againstscores[..., None]=[1, 91, 6, 1]and aborted. Note12288 = 6 * 2048and91 = 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.rsSwitchGLU::forward: drop the rank-assuming final reshape and return the already-squeezed[..., top_k, hidden], matching the provendeepseek_v2sibling and the sharedmoe_weighted_sumcontract. 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_v32has its own already-correct SwitchGLU and is unaffected.switch_glu_forward_preserves_rank_for_2d_and_3dandmoe_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