model, mtmd: fix gemma4 vision handling (#28335)

* model, mtmd: fix gemma4 vision handling

* nits

src/llama-hparams.h

@@ -28,6 +28,14 @@ enum llama_swa_type {
     LLAMA_SWA_TYPE_SYMMETRIC = 3,
 };
 
+// how the non-causal mask should be constructed with llama_set_causal_attn(ctx, false)
+// (e.g. mtmd decoding image tokens)
+enum llama_non_causal_type {
+    LLAMA_NON_CAUSAL_TYPE_ALL      = 0, // all layers non-causal, SWA still applied (gemma 3, qwen-vl, ...)
+    LLAMA_NON_CAUSAL_TYPE_SWA_ONLY = 1, // SWA layers non-causal, dense layers stay causal (gemma 4)
+    LLAMA_NON_CAUSAL_TYPE_SWA_FULL = 2, // all layers non-causal, SWA not applied between tokens of the current ubatch (deepseek 4)
+};
+
 // forward declaration; full definition in llama-graph.h
 enum llm_ffn_op_type : int;
 
@@ -164,9 +172,9 @@ struct llama_hparams {
     // the size of the sliding window (0 - no SWA)
     uint32_t n_swa = 0;
 
-    // deepseek4 vision: when decoding non-causally (multimodal input), SWA is not applied between tokens of the current ubatch (the image span); older tokens are still window-clipped
-    // for other models (like gemma 3, gemma 4): SWA is always applied to match transformers implementation
-    bool swa_full_non_causal = false;
+    // see llama_non_causal_type
+    // note: for SWA_FULL, older tokens (outside the current ubatch) are still window-clipped
+    llama_non_causal_type non_causal_type = LLAMA_NON_CAUSAL_TYPE_ALL;
 
     // if is_swa_impl[il] == 1, then layer il is SWA
     // if is_swa_impl[il] == 0, then layer il is dense (i.e. non-SWA)

src/llama-kv-cache.cpp

@@ -1681,8 +1681,8 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data
 
                 // apply SWA if any
                 if (swa) {
-                    // see llama_hparams::swa_full_non_causal
-                    const bool in_span = !causal && args.hparams.swa_full_non_causal && p0 >= seq_pos_min[seq_id];
+                    // see llama_non_causal_type
+                    const bool in_span = !causal && args.hparams.non_causal_type == LLAMA_NON_CAUSAL_TYPE_SWA_FULL && p0 >= seq_pos_min[seq_id];
                     if (!in_span && llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1)) {
                         goto skip;
                     }
@@ -1754,6 +1754,12 @@ void llama_kv_cache::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * u
     // n_tps == n_tokens_per_stream
     const int64_t n_tps = n_tokens/n_stream;
 
+    // see llama_non_causal_type
+    // only the SWA cache (or the SWA layers of a single cache) become non-causal
+    if (!causal_attn && hparams.non_causal_type == LLAMA_NON_CAUSAL_TYPE_SWA_ONLY) {
+        causal_attn = swa_type == LLAMA_SWA_TYPE_NONE;
+    }
+
     //const int64_t t_start = ggml_time_us();
 
     const args_set_input_kq_mask args = {

src/llama-model.cpp

@@ -1956,6 +1956,7 @@ void llama_model::print_info() const {
         LLAMA_LOG_INFO("%s: n_rot                 = %u\n",     __func__, hparams.n_rot_full);
         LLAMA_LOG_INFO("%s: n_swa                 = %u\n",     __func__, hparams.n_swa);
         LLAMA_LOG_INFO("%s: is_swa_any            = %u\n",     __func__, hparams.is_swa_any());
+        LLAMA_LOG_INFO("%s: non_causal_type       = %d\n",     __func__, hparams.non_causal_type);
         LLAMA_LOG_INFO("%s: n_embd_head_k         = %u\n",     __func__, hparams.n_embd_head_k_full);
         LLAMA_LOG_INFO("%s: n_embd_head_v         = %u\n",     __func__, hparams.n_embd_head_v_full);
         LLAMA_LOG_INFO("%s: n_gqa                 = %s\n",     __func__, print_f([&](uint32_t il) { return hparams.n_gqa(il);        }, hparams.n_layer_all).c_str());

src/models/deepseek4.cpp

@@ -68,7 +68,7 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
     hparams.set_swa_pattern(0);
     // tokens of an image span attend bidirectionally to the whole span, the window only applies to older tokens
     // ref: get_window_topk_idxs_visible in the reference impl
-    hparams.swa_full_non_causal = true;
+    hparams.non_causal_type = LLAMA_NON_CAUSAL_TYPE_SWA_FULL;
     for (uint32_t il = hparams.n_layer(); il < hparams.n_layer_all; ++il) {
         hparams.is_swa_impl[il] = true;
     }

src/models/gemma4.cpp

@@ -19,6 +19,11 @@ void llama_model_gemma4::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_SWA,  hparams.n_embd_head_v_swa);
     ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING,     hparams.f_final_logit_softcapping, false);
 
+    // when non_causal is set, the model will use bidirectional attention on SWA layers only, while dense layers will remain causal
+    // ref: use_bidirectional_attention == "vision" in HF config
+    // note: E2B/E4B are always causal, bypassing this logic
+    hparams.non_causal_type = LLAMA_NON_CAUSAL_TYPE_SWA_ONLY;
+
     switch (hparams.n_layer()) {
         case 30: type = LLM_TYPE_26B_A4B; break;
         case 35: type = LLM_TYPE_E2B; break;

tools/mtmd/clip.cpp

@@ -1636,8 +1636,7 @@ struct clip_model_loader {
                             hparams.patch_size = hparams.patch_size * hparams.n_merge;
                             hparams.n_merge = 1;
                         }
-                        // @ngxson : the model performs quite poor with small images, we need to bump minimum image tokens to 40 to avoid that
-                        hparams.set_limit_image_tokens(40, 280);
+                        hparams.set_limit_image_tokens(70, 1120);
                         hparams.set_warmup_n_tokens(256); // avoid OOM on warmup
                     } break;
 

tools/mtmd/mtmd.cpp

@@ -2173,9 +2173,11 @@ bool mtmd_decode_use_non_causal(const mtmd_context * ctx, const mtmd_input_chunk
         proj_type = ctx->proj_type_a();
     }
     switch (proj_type) {
-        case PROJECTOR_TYPE_GEMMA3:
         case PROJECTOR_TYPE_GEMMA4V:
+            // E2B (n_embd = 1536) and E4B (n_embd = 2560) always use causal
+            return ctx->n_embd_text != 1536 && ctx->n_embd_text != 2560;
         case PROJECTOR_TYPE_GEMMA4UV:
+        case PROJECTOR_TYPE_GEMMA3:
         case PROJECTOR_TYPE_DEEPSEEK4V:
             return true;
         default: