fix(deepseek_v32): add causal-mask fallback so prefill is not bidirectional (same class as #618)

[inureyes]

Summary

src/models/deepseek_v32.rs (DeepSeek-V3.2, the DSA lightning-indexer model added in #583) has the same latent non-causal-prefill bug that #618 just fixed in the shared DeepSeek-V3 backbone (src/models/deepseek_v3.rs). In MLAAttention::forward, causality is applied to the attention scores only when the caller supplies a mask; there is no causal fallback for mask == None. Because both standard generation paths pass mask == None for prefill (the offline CLI text prefill, and the VLM embeddings prefill where merge_llava intentionally returns no mask), a multi-token prefill attends bidirectionally and writes future-contaminated K/V into the cache.

Current code (MLAAttention::forward, around lines 403-408):

// Apply causal mask to pe_scores
let pe_scores = if let Some(m) = mask {
    mlxcel_core::add(&pe_scores, m)
} else {
    pe_scores   // no causal fallback: prefill is bidirectional when mask == None
};

Contrast with the fixed deepseek_v3.rs (post-#618), which adds the l > 1 fallback:

let pe_scores = if let Some(m) = mask {
    mlxcel_core::add(&pe_scores, m)
} else if l > 1 {
    let causal = create_causal_mask(l, live_before);
    mlxcel_core::add(&pe_scores, &causal)
} else {
    // decode (l == 1): every cached position is causally valid, no mask
    pe_scores
};

Why this matters

pe_scores is fed as the additive SDPA mask to attention_from_ptr in the prefill branch (l > 1), so causality exists only if it is baked into pe_scores. The C++ fast_scaled_dot_product_attention wrapper applies no implicit causality for an array mask. When mask == None, a multi-token prefill therefore attends to future tokens, and every decoder layer above the first writes future-contaminated K/V into the cache. That corrupts the whole generation (not just the first token), which is exactly the failure class #618 root-caused for Kimi-VL on the deepseek_v3 backbone: degraded or garbage output, e.g. a constant repeated token. Any real DeepSeek-V3.2 checkpoint run through the CLI text path or the VLM embeddings path is exposed.

There is a second, DSA-specific facet. The sparse-prefill helper apply_sparse_prefill_mask documents its own invariant as "pe_scores already carries the causal mask, so keeping it where selected reproduces upstream's sparse_mask & mask" (around lines 520-523). That invariant is false when mask == None, so the sparse path silently drops causality too. In addition, the lightning indexer's top-k selection idx.top_indices(idx_q, idx_k, idx_w, mask) (around line 392) receives the same None, so top-k key selection during prefill can pick future positions. The fix must ensure causality reaches both the dense pe_scores additive mask and the indexer top-k selection.

Where

• src/models/deepseek_v32.rs, MLAAttention::forward: the pe_scores masking at lines ~403-408; prefill SDPA at lines ~459-483; sparse-prefill helper and its doc comment at ~520-546; indexer top-k at ~392. The RoPE offset is captured at line ~339 (let offset = cache.offset;) before cache.update_and_fetch at ~376, so the fallback needs the pre-update live length, mirroring deepseek_v3's live_before.
• Proven fix pattern: src/models/deepseek_v3.rs DeepSeekV3Attention::forward (post-fix(models/deepseek-v3): restore last layer, causal prefill, f16 clip #618) and its prefill_is_causal_without_caller_mask test.
• Reference fallbacks already in the tree: src/models/qwen3.rs (else if l > 1 && mask.is_none(), around line 200) and src/models/deepseek_v2.rs.

What to check / do

• Add the causal-mask fallback in MLAAttention::forward: when mask.is_none() && l > 1, add create_causal_mask(l, live_len_before) to pe_scores, capturing the live length before cache.update_and_fetch (mirroring deepseek_v3's live_before). Use create_causal_mask from mlxcel_core::utils, as deepseek_v3 does.
• Confirm the fallback also constrains the sparse path: the additive causal mask must be present in pe_scores before apply_sparse_prefill_mask runs so its stated invariant holds, and verify the indexer top_indices does not select future keys during a maskless multi-token prefill (thread causality into the indexer path if it does).
• Verify the decode path (l == 1) is unaffected: decode is always called with mask == None and every cached or gathered position is causally valid, so it must NOT get a causal mask.
• Sanity-check the decoder layer count against num_hidden_layers. Note: on inspection deepseek_v32's sanitize_weights strips only the out-of-range index (mtp_layer = num_hidden_layers, keeping layer_idx < num_hidden_layers) and from_weights builds 0..num_hidden_layers, so it does NOT appear to share the deepseek_v3 off-by-one that dropped the final real decoder layer. Add a regression test to lock this in rather than assuming it.

Acceptance Criteria

• A multi-token prefill with mask == None is causal: position 0 of an N-token prefill produces the same result as a single-token forward of that same token (the assertion used by deepseek_v3's prefill_is_causal_without_caller_mask).
• A unit test reproduces the bug (fails before the fix, passes after) for deepseek_v32 attention, covering both the dense and the sparse-indexer prefill paths.
• The absorbed decode path (l == 1) output is unchanged by the fix (regression-checked against the pre-fix decode result or the materialized prefill path, as deepseek_v3's test cross-checks).
• A regression test asserts the model builds all num_hidden_layers decoder blocks and strips only the MTP trailer at index num_hidden_layers.
• Real-checkpoint generation on a genuine DeepSeek-V3.2 checkpoint is coherent (no degraded or constant-token output) through the CLI text path and, where applicable, the VLM embeddings path.

References

• PR fix(models/deepseek-v3): restore last layer, causal prefill, f16 clip #618 (fix(models/deepseek-v3): restore last layer, causal prefill, f16 clip) root-caused and fixed this exact class on the deepseek_v3 backbone and explicitly flagged deepseek_v32 as carrying the same "causality only via caller mask" structure. Follow its fix pattern and the prefill_is_causal_without_caller_mask test.
• feat(deepseek-v3.2 / glm-moe-dsa): implement the DSA lightning indexer #583 added DeepSeek-V3.2 (the DSA lightning indexer).