[Bug] Omni inherits embedding attention masks, breaking Gemma 4 mixed-attention generation

Summary

Gemma 4 E4B loads successfully but produces no tokens on the CPU backend at 2edeef91b425e98a93707840b6fffdd97980bdbe. Both text-only and image prompts fail during the first forward pass.

The cause appears to be the attention-mask virtual dispatch after Omni began inheriting from Embedding: generation uses Embedding::gen_attention_mask() instead of Llm::gen_attention_mask(). For attention_type=mix, this supplies a scalar mask where the model expects the full/sliding mask tensor.

Platform and version

• Host: Windows; cross-compiled for HarmonyOS arm64-v8a.
• Device: HUAWEI MatePad Edge (QXS-W20), 32 GB RAM, CPU backend.
• SDK native toolchain: 6.1.1.125, Clang 15; OHOS platform level 9, c++_shared.
• Failing revision: 2edeef91b425e98a93707840b6fffdd97980bdbe.
• Working baseline: MNN 3.6.0, cc20f672af9e177e2fa338c332dc097de2fc9264.
• I also checked master 47ccf6c6bb5b6d357cd1f9b4370cbecb6188fd34: Omni still inherits Embedding without overriding this method. This newer commit was inspected, not device-tested.

Build: Release shared library, MNN_BUILD_LLM=ON, MNN_BUILD_LLM_OMNI=ON, MNN_LOW_MEMORY=ON, MNN_SUPPORT_TRANSFORMER_FUSE=ON, MNN_ARM82=ON, MNN_SME2=OFF, OpenCV/image codecs enabled, OpenCL disabled. SME2 was disabled because of the SDK compiler's intrinsic support. The library and signed application build successfully; this is a runtime failure.

Model and reproduction

Model: MNN/gemma-4-E4B-it-MNN, using the same local model files for all comparisons, without re-exporting or modifying them.

The tested llm.mnn SHA-256 is b55bad1efd4e66217e7ff75b896535fc01a0d07fcc78cc7a3d6b562fd71e8836.

Relevant llm_config.json fields:

{
  "model_type": "gemma4",
  "attention_mask": "float",
  "attention_type": "mix",
  "is_visual": true,
  "is_mrope": false
}

Reproduced through a native HarmonyOS app using MNN's C++ LLM API, not through llm_demo:

1. Load the model with the CPU backend.
2. Send a text-only prompt: 计算12加8，只输出结果。 (calculate 12 + 8).
3. The first forward fails and returns no generated tokens. Releasing/reloading the model does not resolve it.
4. A prompt with an image fails at the same node. An image is therefore not required to reproduce the failure.

Runtime logs and diagnosis

Compute Shape Error for /Gather_5_output_0
code=3 in onForward, 779
[Error]: onForward returned no outputs. seqLen=87, inDecode=0, inputs=5, moduleKey=(100,0)

Temporary diagnostic logging showed:

attention_type=mix, mask rank=0
Gather /Gather_5_output_0: input count=3, params rank=0, axis=0, indices rank=1

Reading the model graph confirmed that this Gather consumes attention_mask. Llm::gen_attention_mask() has a mixed-attention path producing {2, 1, 1, seq_len, kv_seq_len}, but the inherited Embedding implementation returns a scalar float zero on CPU. The Gather cannot index that scalar.

Relevant code at the tested revision:

• Omni inheritance
• Embedding mask implementation
• LLM mixed-attention mask implementation

Locally verified fix

Declare gen_attention_mask(int seq_len) override in Omni and dispatch according to its existing embedding-mode flag:

VARP Omni::gen_attention_mask(int seq_len) {
    return mIsEmbedding ? Embedding::gen_attention_mask(seq_len)
                        : Llm::gen_attention_mask(seq_len);
}

Same device and model files:

Qwen3-VL-8B image inference and MiniCPM5-2B text inference also passed with the patched library. Embedding-mode inference has not been device-tested; the proposed dispatch retains its existing implementation. No model-name-specific workaround or change to Gather validation was used.

I checked the existing Gemma reports, including #4835 (Windows CPU/Vulkan crashes) and #4641 (Android E2B image crash). They may be useful context, but I could not establish that they share this root cause. This report concerns the specifically diagnosed mixed-mask dispatch regression above.