[inureyes]

Summary

Implements ModelType::EmbeddingGemma2 to support Google's EmbeddingGemma 2 (google/embeddinggemma-2 and mlx-community/embeddinggemma-2-* in bf16, 8bit, 6bit, 5bit, 4bit, mxfp8, mxfp4, nvfp4), a 0.6B Gemma 4-based multimodal embedding architecture (architectures: ["EmbeddingGemma2Model"], model_type: "embedding_gemma2").

Key Changes

• Config parsing & mixed quantization (src/models/embedding_gemma2_config.rs):
  • Parses both top-level multimodal configs (model_type = "embedding_gemma2") and standalone text configs (model_type = "embedding_gemma2_text"), enforcing positive lower bounds (.max(1)) on geometry dimensions.
  • Resolves per_layer_config overrides on zero-padded ("05", "11", "17", "23") and unpadded keys (head_dim = 512, num_key_value_heads = 1, rope_theta = 1_000_000.0, attention_type = "full_attention"), unconditionally defaulting non-sliding layers to global geometry when specific fields are omitted.
  • Handles mlx-vlm mixed quantization overrides ("language_model.ple": false boolean overrides and per-module bit/group dicts across both text_config and root config) via quant_params_for_prefix.
  • Normalizes both Hugging Face (model.language_model.layers.*) and mlx-vlm (language_model.layers.*, language_model.ple.*, language_model.embedding_projection.*) checkpoint weight layouts via sanitize_embedding_gemma2_weights.
• Precision preservation (src/models/sanitize.rs, src/models/detection.rs):
  • Extends bf16_to_f16_at_load and is_embedding_checkpoint to inspect top-level and nested text_config model_type / architectures so embedding_gemma2 / embedding_gemma2_text / EmbeddingGemma2Model preserve BFloat16 weights at load time and avoid FP16 overflow on unscaled pre-norm activations.
• Per-layer attention, MLP, and PLE blocks (src/models/embedding_gemma2_layers.rs):
  • EmbeddingGemma2Attention: per-layer head_dim, num_key_value_heads, and rope_theta; q_norm and k_norm with learnable scale (RMSNorm) and parameterless v_norm (RMSNormNoScale via fast_rms_norm_no_weight); bidirectional RoPE and SDPA with scale = 1.0.
  • EmbeddingGemma2Mlp: GeGLU with gelu_pytorch_tanh (compiled_geglu_approx_activation).
  • EmbeddingGemma2PleBlock & EmbeddingGemma2Layer: gated per-layer residual block (per_layer_input_gate -> GeLU -> * per_layer_input -> per_layer_projection -> post_per_layer_input_norm) and learnable layer_scalar (mlxcel_core::multiply_scalar).
• Multimodal & text forward pass (src/models/embedding_gemma2.rs):
  • sqrt(hidden_size) input embedding scaling and projection-only PLE (language_model.ple.per_layer_model_projection scaled by hidden_size^-0.5 + per_layer_projection_norm across 24 layers).
  • Optional Gemma4VisionModel + Gemma4MultimodalEmbedder (<|image|> 258880) and AudioEncoder + Gemma4MultimodalEmbedder (<|audio|> 258881) placeholder replacement via masked_scatter, with guards against empty image lists or zero-length audio sequences.
  • Bidirectional full-attention mask (create_bidirectional_padding_mask) and bidirectional sliding-window mask (create_bidirectional_window_mask(attention_mask, sliding_window + 1)), bypassing $O(L^2)$ mask allocation when $L \le \text{sliding_window} + 1$.
  • In-model language_model.embedding_projection (512 -> 768, bias=False) before masked mean pooling (PoolingMode::Mean), supporting Matryoshka truncation (768 / 512 / 384 / 256 / 128) and L2 normalization through EmbeddingEngine.
• Registration & Documentation:
  • Wired ModelType::EmbeddingGemma2 into src/models/mod.rs, src/models/detection.rs, src/embeddings/loader.rs, src/model_metadata.rs, src/models/registry.rs, src/execution/memory_estimate.rs, src/distributed/tensor_parallel/inference.rs, docs/embeddings.md, docs/supported-models.md, and TECHNICAL_REPORTS/2316-embedding-gemma2-embedder-20261015.{en,ko}.md.

Changes during review

• Replaced placeholder string in format_text with clean identity formatting and preserved leading tokens in expand_image_tokens.
• Fixed non-sliding layer fallback geometry in EmbeddingGemma2Config::from_json and added boolean false mixed quantization override handling in quant_params_for_prefix.
• Added empty-list / zero-length audio guards in forward_multimodal and bypassed $O(L^2)$ sliding-window mask allocation when $L \le \text{sliding_window} + 1$.

Test plan

• Added 8 unit tests and 1 ignored real-checkpoint parity test in src/models/embedding_gemma2_tests.rs:
  • per_layer_config_resolves_local_and_global_layers
  • config_bounds_guard_against_malformed_nonpositive_values
  • detection_accepts_embedding_gemma2_variants_and_preserves_bf16
  • sanitize_and_quant_params_handle_hf_and_mlx_layouts
  • synthetic_forward_padding_invariance_bidirectionality_and_matryoshka
  • sliding_window_mask_blocks_tokens_beyond_window
  • short_sequence_sliding_window_bypass_matches_explicit_window_mask
  • expand_image_tokens_text_only_and_vision_paths
  • real_checkpoint_embeddinggemma2_parity_and_ranking (#[ignore], checks retrieval_scores_venus_mars: [0.6775, 0.8348])
• Verified cargo test --lib embedding_gemma2, cargo test --lib embeddings::, cargo test --lib models::detection_tests, cargo clippy --lib --tests -- -D warnings, and cargo fmt --check.

Closes #2314

[inureyes]

Implementation Review Summary

Intent

Support Google's EmbeddingGemma 2 (embedding_gemma2) multimodal bidirectional embedder across text, vision, and audio.

Findings Addressed

• Remove corrupted placeholder string <|image><|image|><image|> in EmbeddingGemma2Model::format_text and preserve leading tokens when expanding image placeholders for empty-text prompts (HIGH)
• Correct fallback geometry in EmbeddingGemma2Config::from_json so non-sliding layers unconditionally use fallback_global_kv_heads (1) and fallback_global_head_dim (512) when omitted from per_layer_config (HIGH)
• Align real_checkpoint_embeddinggemma2_parity_and_ranking integration test with official Venus/Mars benchmark queries and reference scores retrieval_scores_venus_mars: [0.6775, 0.8348] (MEDIUM)
• Add unit test expand_image_tokens_text_only_and_vision_paths covering multimodal placeholder expansion (MEDIUM)
• Fallback to text_config.model_type during embedding checkpoint detection in src/models/detection.rs (MEDIUM)
• Add bilingual technical review reports under TECHNICAL_REPORTS/ (MEDIUM)

Remaining Items

None

Verification

• All stated requirements implemented
• No placeholder/mock code remaining
• Integrated into project code flow
• Project conventions followed
• Existing modules reused where applicable
• No unintended structural changes
• Tests pass

[inureyes]

Security & Performance Review Summary

Review Scope

Comprehensive security and performance audit of PR #2316 (embedding_gemma2 multimodal embedder).

Findings Addressed

• Prevent server panic / denial of service when multimodal preprocessing outputs 0 feature tensors or null pixel values in concat_many (HIGH)
• Prevent divide-by-zero panic in masked_scatter index remainder calculation when processing zero-length audio sequences (HIGH)
• Support boolean false overrides in mixed quantization configurations, correctly resolving (0, 0) unquantized parameters for modules such as ple / language_model.ple (HIGH)
• Eliminate O(L^2) quadratic memory allocation for bidirectional sliding window masks when input length is within sliding_window + 1 tokens (HIGH)
• Eliminate in-loop host slice allocations and redundant dtype conversions in layer_scalar, replacing them with mlxcel_core::multiply_scalar (MEDIUM)
• Eliminate redundant cloning of inputs_embeds in forward_from_embeds (MEDIUM)
• Enforce positive lower bounds (.max(1)) on configuration geometry to guard against divide-by-zero panics from malformed JSON (LOW)
• Support text_config.architectures alongside root architectures in model detection and weight sanitization, stripping model prefixes before discarding lm_head.* (LOW)
• Update bilingual technical reports (TECHNICAL_REPORTS/pr-2316-embedding-gemma2-embedder.{en,ko}.md) with review findings and resolutions (LOW)

Verification

• cargo test --lib embedding_gemma2: 8 passed, 0 failed, 1 ignored (real checkpoint test)
• cargo test --lib embeddings::: 85 passed, 0 failed
• cargo clippy --lib --tests -- -D warnings: 0 warnings

[inureyes]

PR Finalization Complete

Summary

• Tests: Verified test coverage across src/models/embedding_gemma2_tests.rs (8 passed, 1 ignored external checkpoint integration test) and embeddings:: (85 passed, 0 failed), covering per-layer attention overrides, projection-only PLE, bidirectional and sliding-window masks, short-sequence bypass, weight sanitization, mixed quantization overrides, and multimodal expansion.
• Documentation: Verified completeness in docs/embeddings.md and docs/supported-models.md, and updated bilingual technical reports (TECHNICAL_REPORTS/pr-2316-embedding-gemma2-embedder.en.md and pr-2316-embedding-gemma2-embedder.ko.md) with verification records.
• Lint & Format: Cleaned and validated formatting across all modified files with cargo fmt --check (commit b51a34fd), and confirmed clean zero-warning status with cargo clippy --lib --tests -- -D warnings.

All checks passing. Ready for merge.