diff --git a/src/llama-hparams.h b/src/llama-hparams.h index e9029ff34b..873399aaa4 100644 --- a/src/llama-hparams.h +++ b/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) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index f22054c3d6..a342ee1191 100644 --- a/src/llama-kv-cache.cpp +++ b/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 = { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index f22e35ed93..9e2e0c8936 100644 --- a/src/llama-model.cpp +++ b/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()); diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 5bdf14b486..6bf9d34449 100644 --- a/src/models/deepseek4.cpp +++ b/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; } diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index c6dd7d1bf9..0cd95742d1 100644 --- a/src/models/gemma4.cpp +++ b/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; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index f2e4875346..74f4e2b5a4 100644 --- a/tools/mtmd/clip.cpp +++ b/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; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index e7f5f114ec..00ecadcf4d 100644 --- a/tools/mtmd/mtmd.cpp +++ b/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: