fix(models): internlm3, hunyuan dense/moe and gemma2 run a fully bidirectional prefill #999

Problem / Background

Four model implementations carry the same inverted attention-mask policy that #991 fixed in src/models/deepseek_v2.rs.

Every generation path (CLI text prefill, the VLM embeddings prefill, chunked prefill) calls LanguageModel::forward with mask == None. The documented contract is that the model then builds its own causal mask: see the doc comment on pad_tokens_for_prefill in src/lib/mlxcel-core/src/generate.rs (line 162), which states that None means "the forward pass will use its built-in causal mask".

These four files instead hand that None straight to an unmasked SDPA. Their multi-token prefill is therefore fully bidirectional: every layer writes future-contaminated K/V into the cache, and generation is wrong from the first sampled token.

Affected files

1. src/models/internlm3.rs

Leaf attention around lines 230-239, model forward InternLM3Model::forward around 411-434. The leaf has the shape:

if l > 1 {
    let mask_ptr = mask.map(...).unwrap_or(std::ptr::null());
    attention_from_ptr(...)
} else {
    causal_attention(...)
}

Causality is applied only on the l == 1 decode side, where it is vacuous, and InternLM3Model::forward never builds a mask.

Observed: ./target/release/mlxcel generate -m models/internlm3-8b-4bit -p "<a 36-token prompt>" -n 12 --no-chat-template degenerates to waswaswaswas....

2. src/models/hunyuan_v1_dense.rs

Leaf around lines 178-187, HunyuanV1DenseModel::forward around 373-398. Identical inverted branch with a pass-through model forward.

Observed: models/hunyuan-1.8b-4bit on the same prompt degenerates to was was was was....

3. src/models/gemma2.rs

Leaf Attention::forward around lines 132-190, unmasked SDPA around 171-183. Callers: Gemma2Model::forward around 389-422 and forward_with_embeddings_impl around 444-481.

This one is worse than the others: the file contains no causal-mask machinery at all. There is no create_causal_mask, no causal_attention, and no mask.is_none() guard anywhere in it. Because Gemma 2 sets attn_logit_softcapping, the null mask reaches compiled_softcap_sdpa in attention_dispatch (src/lib/mlxcel-core/src/layers.rs), which applies no causality either.

The module is reached two ways: standalone gemma2 checkpoints (Gemma2Model::load is the dir_loader for ModelType::Gemma2 in src/model_metadata.rs) and the PaliGemma 2 VLM path in src/loading/vlm_siglip.rs.

Observed: models/gemma2-2b-4bit on a 36-token prompt diverges from the mlx-lm 0.31.3 greedy reference at the very first generated token (mlxcel emits " to", the reference emits " the") while the output stays fluent. That is the mild-degradation form of the same defect. A short prompt such as "The capital of France is" still answers correctly, which is why this went unnoticed.

4. src/models/hunyuan_moe.rs

Leaf around lines 687-696, HunyuanMoeModel::forward around 862-893. Same inverted branch and pass-through model forward.

Not run locally. models/hunyuan-a13b-instruct-4bit is the only local hunyuan checkpoint, and its bf16 sibling exceeds the development machine's memory. This family rests on code inspection only.

Proposed Solution

Follow the convention already used by src/models/llama3.rs (search for else if l > 1 && mask.is_none()): route a maskless multi-token prefill through mlxcel_core::causal_attention, and use attention_from_ptr only when a real mask pointer exists. src/models/deepseek.rs and src/models/dots1.rs do the same.

An alternative that also works, used by src/models/deepseek_v3.rs and src/models/minicpm3.rs, is to build create_causal_mask(seq_len, offset) in the model forward when the caller passed none.

causal_attention keeps decode byte-identical, because it takes a single-query maskless fast path at l == 1.

Scope note

src/models/minicpm3.rs has the risky-looking leaf shape but is SAFE: MiniCPM3Model::forward overrides with create_causal_mask(seq_len, offset) whenever seq_len > 1. Do not "fix" it.

A separate, smaller observation about that same file: its override discards a caller-supplied padded-prefill mask. That is a different concern and is out of scope here.

Acceptance Criteria

• Each of the four files (internlm3.rs, hunyuan_v1_dense.rs, gemma2.rs, hunyuan_moe.rs) routes a maskless multi-token prefill through a causal path, matching the llama3.rs convention.
• models/internlm3-8b-4bit, models/hunyuan-1.8b-4bit and models/gemma2-2b-4bit generate coherently and match an mlx-lm greedy reference on a prompt long enough to expose the defect. A 6-token prompt is not long enough; use 30 or more tokens.
• models/internlm3-8b-4bit needs trust_remote_code=True to load under the mlx-lm reference harness. Capture its reference accordingly, or gate that family on a coherence assertion instead of reference ids.
• A real-checkpoint regression gate exists for at least the families with a local checkpoint, in the style of tests/deepseek_v2_greedy_parity.rs (added by fix(models): deepseek_v2 generates degenerate output on a real checkpoint while mlx-lm is correct on the same weights #991), which pairs a reference-id assertion with a reference-free assertion that adding later prompt tokens does not move the first position's logits. The second form pins the causality property directly and is portable across families.
• src/models/hunyuan_moe.rs is fixed by inspection, and the absence of a runnable checkpoint is stated explicitly in the PR rather than implied to be covered.
• The PaliGemma 2 VLM path (src/loading/vlm_siglip.rs) is re-verified after the gemma2.rs change, since it shares the module.

Technical Considerations

gemma2.rs needs more than a branch flip: the file has no mask helper imported at all, so the causal path has to be introduced rather than re-routed. Its softcapping dispatch (compiled_softcap_sdpa) must receive the mask, otherwise the fix is a no-op for Gemma 2 specifically.

Refs: #991, #998, #525.