fix(deepseek_v32): add causal-mask fallback so maskless prefill is not bidirectional#667
Merged
Merged
Conversation
…t bidirectional MLAAttention::forward applied causality only when the caller supplied a mask, and both standard generation paths pass mask == None for prefill (the offline CLI text prefill and the VLM embeddings prefill), so a multi-token prefill attended bidirectionally and wrote future-contaminated K/V into the cache; the DSA facets were also exposed (the lightning indexer's top-k could select future keys, and apply_sparse_prefill_mask's "pe_scores already carries the causal mask" invariant was false). This is the same class PR #618 fixed on the deepseek_v3 backbone and explicitly flagged here. The fix mirrors deepseek_v3: capture the pre-update live length (live_before) next to the rope offset, build create_causal_mask(l, live_before) once for any maskless multi-token forward, and route that effective mask to BOTH consumers, the indexer top_indices selection and the additive pe_scores mask; decode (l == 1) stays maskless since every cached position is causally valid. Four regression tests: the dense prefill causality test fails before the fix (position 0 contaminated by position 1, max diff 9e-3) and passes after; the sparse-indexer variant pins the kv_len > index_topk path; a decode-after-prefill test pins the maskless decode step; a sanitize test locks that only the MTP trailer at layer_idx == num_hidden_layers is stripped (the #618-class off-by-one). Real-checkpoint validation is not feasible on the 128 GB validation machine: every public deepseek_v32-backbone checkpoint (DeepSeek-V3.2, GLM-5 via glm_moe_dsa) exceeds it; the fix is a byte-for-byte mirror of the #618 pattern that was validated on real hardware via the deepseek_v3 backbone, with the fail-then-pass unit proof standing in here.
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.
Summary
MLAAttention::forwardinsrc/models/deepseek_v32.rsapplied causality to the attention scores only when the caller supplied a mask. Both standard generation paths passmask == Nonefor prefill (the offline CLI text prefill and the VLM embeddings prefill), so a multi-token prefill attended bidirectionally and every layer wrote future-contaminated K/V into the cache. This is the class PR #618 root-caused and fixed on thedeepseek_v3backbone, which explicitly flaggeddeepseek_v32as carrying the same structure. Two DSA-specific facets were exposed as well: the lightning indexer'stop_indicesreceived the sameNone(top-k selection could pick future keys), andapply_sparse_prefill_mask's documented invariant ("pe_scores already carries the causal mask") was false.Fix
Mirrors
deepseek_v3post-#618: capture the pre-update live length (live_before = cache.live_len()) next to the rope offset, buildcreate_causal_mask(l, live_before)once for any maskless multi-token forward, and route that effective mask to BOTH consumers, the indexer top-k selection and the additivepe_scoresmask. Decode (l == 1) stays maskless: every cached or gathered position is causally valid.Validation
kv_len > index_topkprefill path post-fix; the existingindexer_top_indices_respects_causal_masktest pins that the additive mask steers top-k selection.l == 1step; a sanitize test locks that only the MTP trailer atlayer_idx == num_hidden_layersis stripped (the fix(models/deepseek-v3): restore last layer, causal prefill, f16 clip #618-class off-by-one, which deepseek_v32 does not have).cargo fmt, andcargo clippy --all-targets -- -D warningsare clean.deepseek_v32-backbone checkpoint exceeds the 128 GB validation machine (DeepSeek-V3.2 ~350 GB at 4-bit; GLM-5 viaglm_moe_dsa~200 GB at 4-bit), so end-to-end generation could not be re-validated here. The fix is a byte-for-byte mirror of the fix(models/deepseek-v3): restore last layer, causal prefill, f16 clip #618 pattern that WAS validated on real hardware via the deepseek_v3 backbone (Kimi-VL), with the fail-then-pass unit proof standing in for this backbone.Closes #619